wip

Author: Lucas McDonald

DynamoDbEncryption/runtimes/python/src/aws_database_encryption_sdk/encrypted/client.py

@@ -107,7 +107,6 @@ def paginate_query(self, **paginate_query_kwargs):
         except KeyError:
             pass
 
-        # client_query_input = self._maybe_transform_request_to_dynamodb_item(item_key = "Key", **client_query_kwargs)
         transformed_request = self._transformer.query_input_transform(
             QueryInputTransformInput(
                 sdk_input = client_query_kwargs
@@ -118,11 +117,6 @@ def paginate_query(self, **paginate_query_kwargs):
             transformed_request["PaginationConfig"] = pagination_config
         sdk_page_response = self._paginator.paginate(**transformed_request)
 
-        print(f"{sdk_page_response=}")
-        print(f"{sdk_page_response.__dict__=}")
-
-
-
         for page in sdk_page_response:
             # boto3 docs are wrong. NextToken is never returned.
 

DynamoDbEncryption/runtimes/python/src/aws_database_encryption_sdk/encrypted/paginator.py

@@ -214,12 +214,9 @@ def keys(self, keys):
 
     def batch_write_item_request(self, batch_write_item_request):
         if "RequestItems" in batch_write_item_request:
-            batch_write_item_request["RequestItems"] = self.batch_write_item_request_items(batch_write_item_request["RequestItems"])
+            batch_write_item_request["RequestItems"] = self.batch_write_item_items(batch_write_item_request["RequestItems"])
         return batch_write_item_request
 
-    def batch_write_item_request_items(self, request_items):
-        return self.batch_write_item_items(request_items)
-    
     def batch_write_item_items(self, items):
         items_out = {}
         for table_name, table_value in items.items():
@@ -252,4 +249,57 @@ def batch_write_item_item_collection_metrics(self, item_collection_metrics):
             item_collection_metrics_out[table_name] = table_value
         return item_collection_metrics_out
 
-    def 
+    def transact_write_items_request(self, transact_write_items_request):
+        if "TransactItems" in transact_write_items_request:
+            transact_write_items_request["TransactItems"] = self.transact_write_items(transact_write_items_request["TransactItems"])
+        return transact_write_items_request
+    
+    def transact_write_items(self, transact_items):
+        output_transact_items = []
+        for transact_item in transact_items:
+            if "Put" in transact_item:
+                transact_item["Put"]["Item"] = self.item(transact_item["Put"]["Item"])
+                output_transact_items.append(transact_item)
+            elif "Delete" in transact_item:
+                transact_item["Delete"]["Key"] = self.item(transact_item["Delete"]["Key"])
+                output_transact_items.append(transact_item)
+            elif "ConditionCheck" in transact_item:
+                transact_item["ConditionCheck"]["Key"] = self.item(transact_item["ConditionCheck"]["Key"])
+                output_transact_items.append(transact_item)
+            elif "Update" in transact_item:
+                # Update requests are not supported yet.
+                raise NotImplementedError('"update_item" is not yet implemented')
+            else:
+                raise ValueError(f"Unknown transact_write_items method key: {transact_item}")
+        return output_transact_items
+
+    def transact_write_items_response(self, transact_write_items_response):
+        return transact_write_items_response
+    
+    def transact_get_items_request(self, transact_get_items_request):
+        if "TransactItems" in transact_get_items_request:
+            transact_get_items_request["TransactItems"] = self.transact_get_items(transact_get_items_request["TransactItems"])
+        return transact_get_items_request
+    
+    def transact_get_items(self, transact_items):
+        output_transact_items = []
+        for transact_item in transact_items:
+            if "Get" in transact_item:
+                transact_item["Get"]["Key"] = self.item(transact_item["Get"]["Key"])
+                output_transact_items.append(transact_item)
+            else:
+                raise ValueError(f"Unknown transact_get_items method key: {transact_item}")
+        return output_transact_items
+    
+    def transact_get_items_response(self, transact_get_items_response):
+        if "Responses" in transact_get_items_response:
+            transact_get_items_response["Responses"] = self.transact_get_items_responses(transact_get_items_response["Responses"])
+        return transact_get_items_response
+    
+    def transact_get_items_responses(self, responses):
+        responses_out = []
+        for response in responses:
+            if "Item" in response:
+                response["Item"] = self.item(response["Item"])
+            responses_out.append(response)
+        return responses_out

DynamoDbEncryption/runtimes/python/src/aws_database_encryption_sdk/encrypted/table.py

@@ -3,6 +3,9 @@
 
 class ClientShapeToResourceShapeConverter(BotoInterfaceShapeConverter):
 
+    def __init__(self, delete_table_name = True):
+        self.delete_table_name = delete_table_name
+
     def item(self, item):
         return ddb_to_dict(item)
 
@@ -12,22 +15,26 @@ def attribute_value(self, attribute_value):
 
     def put_item_request(self, put_item_request):
         out = super().put_item_request(put_item_request)
-        del out["TableName"]
+        if self.delete_table_name:
+            del out["TableName"]
         return out
 
     def get_item_request(self, get_item_request):
         out = super().get_item_request(get_item_request)
-        del out["TableName"]
+        if self.delete_table_name:
+            del out["TableName"]
         return out
 
     def query_request(self, query_request):
         out = super().query_request(query_request)
-        del out["TableName"]
+        if self.delete_table_name:
+            del out["TableName"]
         return out
 
     def scan_request(self, scan_request):
         out = super().scan_request(scan_request)
-        del out["TableName"]
+        if self.delete_table_name:
+            del out["TableName"]
         return out
 
     def expression(self, condition_expression, expression_attribute_names, expression_attribute_values):

DynamoDbEncryption/runtimes/python/src/aws_database_encryption_sdk/encrypted/utils.py

@@ -85,7 +85,7 @@ def expression(self, condition_expression, expression_attribute_names, expressio
 
     def batch_write_item_request_items(self, tables):
         """Transform a batch write request's items to DynamoDB format."""
-        output_tables = []
+        output_tables = {}
         table_names = list(tables.keys())
         for table_name in table_names:
             requests = tables[table_name]
@@ -105,7 +105,7 @@ def batch_write_item_request_items(self, tables):
                 else:
                     raise ValueError(f"Unknown batch_write_items method key: {request_name}")
                 output_requests.append({request_name: boto3_request})
-            output_tables.append({table_name: output_requests})
+            output_tables[table_name] = output_requests
         return output_tables
 
     def batch_get_request(self, **kwargs):