From bfac8a8a07f694dcbfa773c052dd7bf45d7fde06 Mon Sep 17 00:00:00 2001
From: Otto Jongerius <otto.jongerius@troveplatform.co.nz>
Date: Mon, 27 Nov 2023 11:15:36 -0800
Subject: [PATCH] DynamoDB: put_item() now returns old item for
 ConditionalCheckFailed exceptions (#7068)

---
 moto/dynamodb/models/__init__.py              |  2 +
 moto/dynamodb/models/table.py                 |  9 +++-
 moto/dynamodb/responses.py                    |  4 ++
 .../exceptions/test_dynamodb_exceptions.py    | 42 +++++++++++++++++++
 4 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/moto/dynamodb/models/__init__.py b/moto/dynamodb/models/__init__.py
index fe12880f0..1e6df206a 100644
--- a/moto/dynamodb/models/__init__.py
+++ b/moto/dynamodb/models/__init__.py
@@ -225,6 +225,7 @@ class DynamoDBBackend(BaseBackend):
         expression_attribute_names: Optional[Dict[str, Any]] = None,
         expression_attribute_values: Optional[Dict[str, Any]] = None,
         overwrite: bool = False,
+        return_values_on_condition_check_failure: Optional[str] = None,
     ) -> Item:
         table = self.get_table(table_name)
         return table.put_item(
@@ -234,6 +235,7 @@ class DynamoDBBackend(BaseBackend):
             expression_attribute_names,
             expression_attribute_values,
             overwrite,
+            return_values_on_condition_check_failure,
         )
 
     def get_table_keys_name(
diff --git a/moto/dynamodb/models/table.py b/moto/dynamodb/models/table.py
index 0e211fd35..dd4db4831 100644
--- a/moto/dynamodb/models/table.py
+++ b/moto/dynamodb/models/table.py
@@ -517,6 +517,7 @@ class Table(CloudFormationModel):
         expression_attribute_names: Optional[Dict[str, str]] = None,
         expression_attribute_values: Optional[Dict[str, Any]] = None,
         overwrite: bool = False,
+        return_values_on_condition_check_failure: Optional[str] = None,
     ) -> Item:
         if self.hash_key_attr not in item_attrs.keys():
             raise MockValidationException(
@@ -571,7 +572,13 @@ class Table(CloudFormationModel):
                 expression_attribute_values,
             )
             if not condition_op.expr(current):
-                raise ConditionalCheckFailed
+                if (
+                    return_values_on_condition_check_failure == "ALL_OLD"
+                    and current is not None
+                ):
+                    raise ConditionalCheckFailed(item=current.to_json()["Attributes"])
+                else:
+                    raise ConditionalCheckFailed
 
         if range_value:
             self.items[hash_value][range_value] = item
diff --git a/moto/dynamodb/responses.py b/moto/dynamodb/responses.py
index 1cf70b350..77197ecf7 100644
--- a/moto/dynamodb/responses.py
+++ b/moto/dynamodb/responses.py
@@ -456,6 +456,9 @@ class DynamoHandler(BaseResponse):
         name = self.body["TableName"]
         item = self.body["Item"]
         return_values = self.body.get("ReturnValues", "NONE")
+        return_values_on_condition_check_failure = self.body.get(
+            "ReturnValuesOnConditionCheckFailure"
+        )
 
         if return_values not in ("ALL_OLD", "NONE"):
             raise MockValidationException("Return values set to invalid value")
@@ -498,6 +501,7 @@ class DynamoHandler(BaseResponse):
             expression_attribute_names,
             expression_attribute_values,
             overwrite,
+            return_values_on_condition_check_failure,
         )
 
         item_dict = result.to_json()
diff --git a/tests/test_dynamodb/exceptions/test_dynamodb_exceptions.py b/tests/test_dynamodb/exceptions/test_dynamodb_exceptions.py
index e3a788c3c..ffca7dab2 100644
--- a/tests/test_dynamodb/exceptions/test_dynamodb_exceptions.py
+++ b/tests/test_dynamodb/exceptions/test_dynamodb_exceptions.py
@@ -674,6 +674,48 @@ def test_put_item_empty_set():
     )
 
 
+@mock_dynamodb
+def test_put_item_returns_old_item():
+    dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
+    table = dynamodb.create_table(
+        TableName="test-table",
+        KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}],
+        AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}],
+        BillingMode="PAY_PER_REQUEST",
+    )
+
+    table.put_item(Item={"pk": "foo", "bar": "baz"})
+
+    with pytest.raises(ClientError) as exc:
+        table.put_item(
+            Item={"pk": "foo", "bar": "quuz"},
+            ConditionExpression="attribute_not_exists(pk)",
+        )
+    resp = exc.value.response
+    assert resp["Error"] == {
+        "Message": "The conditional request failed",
+        "Code": "ConditionalCheckFailedException",
+    }
+    assert resp["message"] == "The conditional request failed"
+    assert "Item" not in resp
+
+    table.put_item(Item={"pk": "foo", "bar": "baz"})
+
+    with pytest.raises(ClientError) as exc:
+        table.put_item(
+            Item={"pk": "foo", "bar": "quuz"},
+            ReturnValuesOnConditionCheckFailure="ALL_OLD",
+            ConditionExpression="attribute_not_exists(pk)",
+        )
+    resp = exc.value.response
+    assert resp["Error"] == {
+        "Message": "The conditional request failed",
+        "Code": "ConditionalCheckFailedException",
+    }
+    assert "message" not in resp
+    assert resp["Item"] == {"pk": {"S": "foo"}, "bar": {"S": "baz"}}
+
+
 @mock_dynamodb
 def test_update_expression_with_trailing_comma():
     resource = boto3.resource(service_name="dynamodb", region_name="us-east-1")