refactor: update DynamoDBStore to use update_item for batch operations

Author: 5enxia

libs/langgraph-checkpoint-aws/langgraph_checkpoint_aws/store/dynamodb/base.py

@@ -571,30 +571,32 @@ def _put_item(
         composite_key = self._construct_composite_key(namespace, key)
         current_time = datetime.now(timezone.utc).isoformat()
 
-        # Check if item exists to preserve created_at
-        existing_item = None
-        try:
-            response = self.table.get_item(Key={"PK": composite_key[0], "SK": composite_key[1]})
-            existing_item = response.get("Item")
-        except Exception:
-            pass
-
-        item: dict[str, Any] = {
-            "PK": composite_key[0],
-            "SK": composite_key[1],
-            "value": value,
-            "created_at": existing_item["created_at"] if existing_item else current_time,
-            "updated_at": current_time,
+        # Build update expression to set created_at only if it doesn't exist
+        update_expression = "SET #value = :value, updated_at = :updated_at"
+        expression_attribute_names = {"#value": "value"}
+        expression_attribute_values = {
+            ":value": value,
+            ":updated_at": current_time,
+            ":created_at": current_time,
         }
 
+        # Add created_at with if_not_exists to preserve existing value
+        update_expression += ", created_at = if_not_exists(created_at, :created_at)"
+
         # Add TTL if configured
         if ttl is not None:
             expires_at = self._calculate_expiry(ttl)
             if expires_at:
-                item["expires_at"] = expires_at
+                update_expression += ", expires_at = :expires_at"
+                expression_attribute_values[":expires_at"] = expires_at
 
         try:
-            self.table.put_item(Item=item)
+            self.table.update_item(
+                Key={"PK": composite_key[0], "SK": composite_key[1]},
+                UpdateExpression=update_expression,
+                ExpressionAttributeNames=expression_attribute_names,
+                ExpressionAttributeValues=expression_attribute_values,
+            )
         except Exception as e:
             logger.error(f"Error putting item {namespace}/{key}: {e}")
             raise

libs/langgraph-checkpoint-aws/tests/unit_tests/store/dynamodb/test_dynamodb_store.py

@@ -235,8 +235,7 @@ def test_batch_get_op_not_found(self, dynamodb_store):
 
     def test_batch_put_op(self, dynamodb_store):
         """Test batch PutOp."""
-        dynamodb_store.table.get_item.return_value = {}
-        dynamodb_store.table.put_item.return_value = {}
+        dynamodb_store.table.update_item.return_value = {}
 
         op = PutOp(
             namespace=("users", "123"),
@@ -248,7 +247,7 @@ def test_batch_put_op(self, dynamodb_store):
         result = dynamodb_store._batch_put_op(op)
 
         assert result is None
-        dynamodb_store.table.put_item.assert_called_once()
+        dynamodb_store.table.update_item.assert_called_once()
 
     def test_batch_put_op_delete(self, dynamodb_store):
         """Test batch PutOp for delete."""