Merge pull request #17 from RSS-Engineering/batch_write_to_backend

Batch write to backend

Author: harinath141

README.md

@@ -51,6 +51,8 @@ NOTE: Rule complexity is limited by the querying capabilities of the backend.
 
 `save()` - store the Model instance to the backend
 
+`batch_save()` - store the List of Model instance to the backend
+
 ### DynamoDB
 
 `get(key: Union[Dict, Any])`

pydanticrud/backends/dynamodb.py

@@ -86,7 +86,6 @@ def _to_epoch_float(dt):
     "array": json.dumps,
 }
 
-
 DESERIALIZE_MAP = {
     "number": float,
     "boolean": bool,
@@ -95,6 +94,11 @@ def _to_epoch_float(dt):
 }
 
 
+def chunk_list(lst, size):
+    for i in range(0, len(lst), size):
+        yield lst[i:i + size]
+
+
 def index_definition(index_name, keys, gsi=False):
     schema = {
         "IndexName": index_name,
@@ -112,7 +116,7 @@ def index_definition(index_name, keys, gsi=False):
 
 class DynamoSerializer:
     def __init__(self, schema):
-        self.properties = schema["properties"]
+        self.properties = schema.get("properties")
         self.definitions = schema.get("definitions")
 
     def _get_type_possibilities(self, field_name) -> Set[tuple]:
@@ -409,3 +413,28 @@ def save(self, item, condition: Optional[Rule] = None) -> bool:
 
     def delete(self, key):
         self.get_table().delete_item(Key=self._key_param_to_dict(key))
+
+    def batch_save(self, items: list) -> dict:
+        """
+        This function is to write multiple records in to dynamodb and returns unprocessed records in dict
+        if something gone wrong with the record.Currently, batch_write is not supporting ConditionExpression
+        Refer docs:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/client/batch_write_item.html
+        """
+        # Prepare the batch write requests
+        request_items = {self.table_name: []}
+
+        # chunk list for size limit of 25 items to write using this batch_write operation refer below.
+        # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/dynamodb/client/batch_write_item.html#:~:text=The%20BatchWriteItem%20operation,Data%20Types.
+        for chunk in chunk_list(items, 25):
+            serialized_items = [self.serializer.serialize_record(item.dict(by_alias=True)) for item in chunk]
+            for serialized_item in serialized_items:
+                request_items[self.table_name].append({"PutRequest": {"Item": serialized_item}})
+        try:
+            response = self.dynamodb.batch_write_item(RequestItems=request_items)
+        except ClientError as e:
+            raise e
+        except (ValueError, TypeError, KeyError) as ex:
+            raise ex
+        unprocessed_items = response.get("UnprocessedItems", {})
+        return unprocessed_items

pydanticrud/backends/sqlite.py

@@ -191,3 +191,6 @@ def save(self, item, condition: Optional[Rule] = None) -> bool:
 
     def delete(self, item_key: str):
         self._conn.execute(f"DELETE FROM {self.table_name} WHERE {self.hash_key} = ?;", [item_key])
+
+    def batch_save(self, items: dict):
+        raise NotImplementedError("This functionality is not yet implemented")

pydanticrud/main.py

@@ -73,3 +73,7 @@ def save(self) -> bool:
     @classmethod
     def delete(cls, *args, **kwargs):
         cls.__backend__.delete(*args, **kwargs)
+
+    @classmethod
+    def batch_save(cls, *args, **kwargs):
+        return cls.__backend__.batch_save(*args, **kwargs)

tests/test_dynamodb.py

@@ -5,6 +5,7 @@
 import random
 
 import docker
+from botocore.exceptions import ClientError
 from pydantic import BaseModel as PydanticBaseModel, Field, root_validator, ValidationError
 from pydanticrud import BaseModel, DynamoDbBackend, ConditionCheckFailed
 import pytest
@@ -66,6 +67,7 @@ class ComplexKeyModel(BaseModel):
     category_id: int
     notification_id: str
     thread_id: str
+    body: str = "some random string"
 
     class Config:
         title = "ComplexModelTitle123"
@@ -241,7 +243,7 @@ def complex_query_data(complex_table):
     accounts = [str(uuid4()) for i in range(4)]
 
     data = [
-        complex_model_data_generator(account=accounts[i % 4], **p)
+        complex_model_data_generator(account=accounts[i % 4], body="some random string", **p)
         for i, p in enumerate(presets)
     ]
     for datum in data:
@@ -528,3 +530,29 @@ def test_alias_model_validator_ingest(dynamo):
     data.pop("typ")
     with pytest.raises(ValidationError):
         AliasKeyModel(**data)
+
+
+def test_batch_write(dynamo, complex_table):
+    response = {"UnprocessedItems": {}}
+    data = [ComplexKeyModel.parse_obj(complex_model_data_generator()) for x in range(0, 10)]
+    un_proc = ComplexKeyModel.batch_save(data)
+    assert un_proc == response["UnprocessedItems"]
+    res_get = ComplexKeyModel.get((data[0].account, data[0].sort_date_key))
+    res_query = ComplexKeyModel.query(
+        Rule(f"account == '{data[0].account}' and sort_date_key == '{data[0].sort_date_key}'")
+    )
+    assert res_get == data[0]
+    assert res_query.count == 1
+    assert res_query.records == [data[0]]
+
+
+def test_message_batch_write_client_exception(dynamo, complex_table):
+    data = [
+        ComplexKeyModel.parse_obj(complex_model_data_generator(body="some big string" * 10000))
+        for x in range(0, 2)
+    ]
+    with pytest.raises(ClientError) as exc:
+        ComplexKeyModel.batch_save(data)
+    assert (
+        exc.value.response["Error"]["Message"] == "Item size has exceeded the maximum allowed size"
+    )