DynamoDB: execute_statement() now supports INSERT/UPDATE/DELETE queries

moto/dynamodb/models/__init__.py

@@ -785,8 +785,6 @@ def execute_statement(
         self, statement: str, parameters: List[Dict[str, Any]]
     ) -> List[Dict[str, Any]]:
         """
-        Only SELECT-statements are supported for now.
-
         Pagination is not yet implemented.
 
         Parsing is highly experimental - please raise an issue if you find any bugs.
@@ -799,7 +797,29 @@ def execute_statement(
                 item.to_json()["Attributes"] for item in table.all_items()
             ]
 
-        return partiql.query(statement, source_data, parameters)
+        return_data, updates_per_table = partiql.query(
+            statement, source_data, parameters
+        )
+
+        for table_name, updates in updates_per_table.items():
+            table = self.tables[table_name]
+            for before, after in updates:
+                if after is None and before is not None:
+                    # DELETE
+                    hash_key = DynamoType(before[table.hash_key_attr])
+                    if table.range_key_attr:
+                        range_key = DynamoType(before[table.range_key_attr])
+                    else:
+                        range_key = None
+                    table.delete_item(hash_key, range_key)
+                elif before is None and after is not None:
+                    # CREATE
+                    table.put_item(after)
+                elif before is not None and after is not None:
+                    # UPDATE
+                    table.put_item(after)
+
+        return return_data
 
     def execute_transaction(
         self, statements: List[Dict[str, Any]]

moto/dynamodb/parsing/partiql.py

@@ -1,12 +1,15 @@
-from typing import TYPE_CHECKING, Any, Dict, List
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
 if TYPE_CHECKING:
     from py_partiql_parser import QueryMetadata
 
 
 def query(
     statement: str, source_data: Dict[str, str], parameters: List[Dict[str, Any]]
-) -> List[Dict[str, Any]]:
+) -> Tuple[
+    List[Dict[str, Any]],
+    Dict[str, List[Tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]]],
+]:
     from py_partiql_parser import DynamoDBStatementParser
 
     return DynamoDBStatementParser(source_data).parse(statement, parameters)

setup.cfg

@@ -56,7 +56,7 @@ all =
     openapi-spec-validator>=0.5.0
     pyparsing>=3.0.7
     jsondiff>=1.1.2
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
     aws-xray-sdk!=0.96,>=0.93
     setuptools
     multipart
@@ -71,7 +71,7 @@ proxy =
     openapi-spec-validator>=0.5.0
     pyparsing>=3.0.7
     jsondiff>=1.1.2
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
     aws-xray-sdk!=0.96,>=0.93
     setuptools
     multipart
@@ -86,7 +86,7 @@ server =
     openapi-spec-validator>=0.5.0
     pyparsing>=3.0.7
     jsondiff>=1.1.2
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
     aws-xray-sdk!=0.96,>=0.93
     setuptools
     flask!=2.2.0,!=2.2.1
@@ -121,7 +121,7 @@ cloudformation =
     openapi-spec-validator>=0.5.0
     pyparsing>=3.0.7
     jsondiff>=1.1.2
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
     aws-xray-sdk!=0.96,>=0.93
     setuptools
 cloudfront =
@@ -144,10 +144,10 @@ dms =
 ds = sshpubkeys>=3.1.0
 dynamodb =
     docker>=3.0.0
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
 dynamodbstreams =
     docker>=3.0.0
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
 ebs = sshpubkeys>=3.1.0
 ec2 = sshpubkeys>=3.1.0
 ec2instanceconnect =
@@ -210,15 +210,15 @@ resourcegroupstaggingapi =
     openapi-spec-validator>=0.5.0
     pyparsing>=3.0.7
     jsondiff>=1.1.2
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
 route53 =
 route53resolver = sshpubkeys>=3.1.0
 s3 =
     PyYAML>=5.1
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
 s3crc32c =
     PyYAML>=5.1
-    py-partiql-parser==0.4.2
+    py-partiql-parser==0.5.0
     crc32c
 s3control =
 sagemaker =

tests/test_dynamodb/test_dynamodb_statements.py

@@ -260,3 +260,109 @@ def test_execute_statement_with_all_clauses(table_name=None):
     partiql_statement = f"SELECT pk FROM \"{table_name}\" WHERE (contains(\"NameLower\", 'code') OR contains(\"DescriptionLower\", 'code')) AND Category = 'free' AND Price >= 0 AND Price <= 1 AND FreeTier IS NOT MISSING AND attribute_type(\"FreeTier\", 'N')"
     items = dynamodb_client.execute_statement(Statement=partiql_statement)["Items"]
     assert items == [{"pk": {"S": "0"}}]
+
+
+@pytest.mark.aws_verified
+@dynamodb_aws_verified()
+def test_insert_data(table_name=None):
+    client = boto3.client("dynamodb", "us-east-1")
+    create_items(table_name)
+    resp = client.execute_statement(
+        Statement=f"INSERT INTO \"{table_name}\" value {{'pk': 'msg3'}}"
+    )
+    assert resp["Items"] == []
+
+    items = client.scan(TableName=table_name)["Items"]
+    assert len(items) == 3
+    assert {"pk": {"S": "msg3"}} in items
+
+    # More advanced insertion
+    client.execute_statement(
+        Statement=f"INSERT INTO \"{table_name}\" value {{'pk': 'msg4', 'attr':{{'sth': ['other']}}}}"
+    )
+
+    items = client.scan(TableName=table_name)["Items"]
+    assert len(items) == 4
+    assert {
+        "pk": {"S": "msg4"},
+        "attr": {"M": {"sth": {"L": [{"S": "other"}]}}},
+    } in items
+
+
+@pytest.mark.aws_verified
+@dynamodb_aws_verified()
+def test_update_data(table_name=None):
+    client = boto3.client("dynamodb", "us-east-1")
+    create_items(table_name)
+
+    items = client.scan(TableName=table_name)["Items"]
+    assert item1 in items
+    assert item2 in items  # unchanged
+
+    # Update existing attr
+    client.execute_statement(
+        Statement=f"UPDATE \"{table_name}\" SET body='other' WHERE pk='msg1'"
+    )
+
+    items = client.scan(TableName=table_name)["Items"]
+    assert len(items) == 2
+    updated_item = item1.copy()
+    updated_item["body"] = {"S": "other"}
+    assert updated_item in items
+    assert item2 in items  # unchanged
+
+    # Set new attr
+    client.execute_statement(
+        Statement=f"UPDATE \"{table_name}\" SET new_attr='asdf' WHERE pk='msg1'"
+    )
+
+    items = client.scan(TableName=table_name)["Items"]
+    assert len(items) == 2
+    updated_item["new_attr"] = {"S": "asdf"}
+    assert updated_item in items
+    assert item2 in items
+
+    # Remove attr
+    client.execute_statement(
+        Statement=f"UPDATE \"{table_name}\" REMOVE new_attr WHERE pk='msg1'"
+    )
+
+    items = client.scan(TableName=table_name)["Items"]
+    assert len(items) == 2
+    updated_item.pop("new_attr")
+    assert updated_item in items
+    assert item2 in items
+
+
+@pytest.mark.aws_verified
+@dynamodb_aws_verified()
+def test_delete_data(table_name=None):
+    client = boto3.client("dynamodb", "us-east-1")
+    create_items(table_name)
+
+    client.execute_statement(Statement=f"DELETE FROM \"{table_name}\" WHERE pk='msg1'")
+
+    items = client.scan(TableName=table_name)["Items"]
+    assert items == [item2]
+
+
+@mock_dynamodb
+def test_delete_data__with_sort_key():
+    client = boto3.client("dynamodb", "us-east-1")
+    client.create_table(
+        TableName="test",
+        AttributeDefinitions=[
+            {"AttributeName": "pk", "AttributeType": "S"},
+            {"AttributeName": "sk", "AttributeType": "S"},
+        ],
+        KeySchema=[
+            {"AttributeName": "pk", "KeyType": "HASH"},
+            {"AttributeName": "sk", "KeyType": "RANGE"},
+        ],
+        BillingMode="PAY_PER_REQUEST",
+    )
+    client.put_item(TableName="test", Item={"pk": {"S": "msg"}, "sk": {"S": "sth"}})
+
+    client.execute_statement(Statement="DELETE FROM \"test\" WHERE pk='msg'")
+
+    assert client.scan(TableName="test")["Items"] == []