Added batch_save method to backend.

Author: Harinath

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

@@ -11,6 +11,7 @@
 
 from ..main import IterableResult
 from ..exceptions import DoesNotExist, ConditionCheckFailed
+from ..dynamo_type_serializer import DynamoTypeSerializer
 
 log = logging.getLogger(__name__)
 
@@ -86,7 +87,6 @@ def _to_epoch_float(dt):
     "array": json.dumps,
 }
 
-
 DESERIALIZE_MAP = {
     "number": float,
     "boolean": bool,
@@ -95,6 +95,10 @@ def _to_epoch_float(dt):
 }
 
 
+def chunk_list(lst, size):
+    return [lst[i: i + size] for i in range(0, len(lst), size)]
+
+
 def index_definition(index_name, keys, gsi=False):
     schema = {
         "IndexName": index_name,
@@ -110,9 +114,10 @@ def index_definition(index_name, keys, gsi=False):
     return schema
 
 
-class DynamoSerializer:
+class DynamoSerializer(DynamoTypeSerializer):
     def __init__(self, schema):
-        self.properties = schema["properties"]
+        super().__init__()
+        self.properties = schema.get("properties")
         self.definitions = schema.get("definitions")
 
     def _get_type_possibilities(self, field_name) -> Set[tuple]:
@@ -409,3 +414,32 @@ 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.This will by default try 3 time if any unprocessed
+        records exists in result. 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, 20):
+            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 }#self.serializer.serialize_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 any unprocessed items
+        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):
+        pass

pydanticrud/dynamo_type_serializer.py

@@ -0,0 +1,54 @@
+import decimal
+from typing import Any, Mapping, Union
+
+NoneType = type(None)
+
+class DynamoTypeSerializer:
+    def __init__(self):
+        self._methods = {
+            bool: serialize_bool,
+            decimal.Decimal: serialize_number,
+            dict: self.serialize_mapping,
+            int: serialize_number,
+            list: self.serialize_list,
+            Mapping: self.serialize_mapping,
+            NoneType: serialize_none,
+            str: serialize_str,
+        }
+
+    def serialize_item(self, item) -> dict:
+        return {k: self.serialize(v) for k, v in item.items()}
+
+    def serialize_mapping(self, value: Mapping) -> dict:
+        return {
+            'M': {k: self.serialize(v) for k, v in value.items()}
+        }
+
+    def serialize_list(self, value: Union[list, tuple]) -> dict:
+        return {'L': [self.serialize(item) for item in value]}
+
+    def serialize(self, value: Any):
+        value_type = type(value)
+        try:
+            route = self._methods[value_type]
+        except KeyError:
+            for type_route, route in self._methods.items():
+                if issubclass(value_type, type_route):
+                    self._methods[value_type] = route
+                    return route(value)
+        else:
+            return route(value)
+        raise TypeError(f'The value {value} is not Dynamodb serializable type.')
+
+def serialize_bool(value: bool) -> dict:
+    return {'BOOL': value}
+
+def serialize_number(value: Union[int, float, decimal.Decimal]) -> dict:
+    return {'N': str(value)}
+
+def serialize_none(value: None) -> dict:
+    return {'NULL': True}
+
+def serialize_str(value: str) -> dict:
+    return {'S': value}
+

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_dynamo_type_serializer.py

@@ -0,0 +1,56 @@
+from decimal import Decimal
+
+from pydanticrud.backends.dynamodb import DynamoSerializer
+
+serializer = DynamoSerializer(schema={})
+def test_int():
+    assert serializer.serialize(0) == {'N': '0'}
+    assert serializer.serialize(34) == {'N': '34'}
+
+
+def test_decimal():
+    assert serializer.serialize(Decimal('1.2345')) == {'N': '1.2345'}
+    assert serializer.serialize(Decimal('0.00')) == {'N': '0.00'}
+
+def test_str():
+    assert serializer.serialize('') == {'S': ''}
+    assert serializer.serialize('hello') == {'S': 'hello'}
+    assert serializer.serialize('hello world') == {'S': 'hello world'}
+
+def test_none():
+    return serializer.serialize(None) == {'NULL': True}
+
+def test_bools():
+    assert serializer.serialize(True) == {'BOOL': True}
+    assert serializer.serialize(False) == {'BOOL': False}
+
+
+def test_list():
+    assert serializer.serialize([1, 'abc', None, None, True]) == {
+        'L': [{'N': '1'}, {'S': 'abc'}, {'NULL': True}, {'NULL': True},
+              {'BOOL': True}]
+    }
+
+def test_maps():
+    assert serializer.serialize({
+        'hello': 'world',
+        'another': 123,
+        'loop': {'Working': True}
+    }) == {
+        'M': {
+            'hello': {'S': 'world'},
+            'another': {'N': '123'},
+            'loop': {'M': {'Working': {'BOOL': True}}}
+        }
+    }
+
+def test_item():
+    assert serializer.serialize_item({
+        'Hello': 'World',
+        'another': 123,
+        'loop': {'Working': True}
+    }) == {
+        'Hello': {'S': 'World'},
+        'another': {'N': '123'},
+        'loop': {'M': {'Working': {'BOOL': True}}}
+    }

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"
@@ -396,7 +398,7 @@ def test_query_with_hash_key_complex(dynamo, complex_query_data):
 
 @pytest.mark.parametrize('order', ('asc', 'desc'))
 def test_ordered_query_with_hash_key_complex(dynamo, complex_query_data, order):
-    middle_record = complex_query_data[(len(complex_query_data)//2)]
+    middle_record = complex_query_data[(len(complex_query_data) // 2)]
     res = ComplexKeyModel.query(
         Rule(f"account == '{middle_record['account']}' and sort_date_key >= '{middle_record['sort_date_key']}'"),
         order=order
@@ -414,8 +416,9 @@ def test_ordered_query_with_hash_key_complex(dynamo, complex_query_data, order):
 @pytest.mark.parametrize('order', ('asc', 'desc'))
 def test_pagination_query_with_hash_key_complex(dynamo, complex_query_data, order):
     page_size = 2
-    middle_record = complex_query_data[(len(complex_query_data)//2)]
-    query_rule = Rule(f"account == '{middle_record['account']}' and sort_date_key >= '{middle_record['sort_date_key']}'")
+    middle_record = complex_query_data[(len(complex_query_data) // 2)]
+    query_rule = Rule(
+        f"account == '{middle_record['account']}' and sort_date_key >= '{middle_record['sort_date_key']}'")
     res = ComplexKeyModel.query(query_rule, order=order, limit=page_size)
     res_data = [(m.account, m.sort_date_key) for m in res]
     check_data = sorted([
@@ -432,13 +435,13 @@ def test_pagination_query_with_hash_key_complex(dynamo, complex_query_data, orde
         (m["account"], m["sort_date_key"])
         for m in complex_query_data
         if m["account"] == middle_record['account'] and m["sort_date_key"] >= middle_record['sort_date_key']
-    ], reverse=order == 'desc')[page_size:page_size*2]
+    ], reverse=order == 'desc')[page_size:page_size * 2]
     assert res_data == check_data
 
 
 def test_pagination_query_with_index_complex(dynamo, complex_query_data):
     page_size = 2
-    middle_record = complex_query_data[(len(complex_query_data)//2)]
+    middle_record = complex_query_data[(len(complex_query_data) // 2)]
     query_rule = Rule(f"account == '{middle_record['account']}' and category_id >= {middle_record['category_id']}")
     check_data = ComplexKeyModel.query(query_rule)
     res = ComplexKeyModel.query(query_rule, limit=page_size)
@@ -528,3 +531,28 @@ 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'