statikk 0.0.8__tar.gz → 0.0.10__tar.gz

{statikk-0.0.8 → statikk-0.0.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: statikk
-Version: 0.0.8
+Version: 0.0.10
 Summary: statikk is a single table application (STA) library for DynamoDb.
 Home-page: https://github.com/terinia/statikk
 Author: Balint Biro

{statikk-0.0.8 → statikk-0.0.10}/docs/usage.rst

@@ -295,7 +295,7 @@ Let's take a look at an example:
 
   card = Card(player_id=player.id, tier="EPIC", values={1, 2, 3, 4}, cost=5, name="FooFoo")
   card.save()
-  Card.update().set("tier", "LEGENDARY").delete("values", {1}).add("cost", 4).remove("name").execute()
+  card.update().set("tier", "LEGENDARY").delete("values", {1}).add("cost", 4).remove("name").execute()
   card = Card.get(card.id)
   card.model_dump()
   # {
@@ -305,6 +305,8 @@ Let's take a look at an example:
   #   "values": {2, 3, 4},
   #   "cost": 9,
   #   "name": "Foo" (default value defined on the model)
+  #   "gsi_pk": "<player_id>",
+  #   "gsi_sk": "Card|LEGENDARY"
   # }
 
 Note that you need to call ``execute`` on the update expression to transmit the changes to the database.

{statikk-0.0.8 → statikk-0.0.10}/src/statikk/engine.py

@@ -6,13 +6,14 @@ import boto3
 from botocore.config import Config
 from pydantic.fields import FieldInfo
 from boto3.dynamodb.conditions import ComparisonCondition, Key
-from boto3.dynamodb.types import TypeDeserializer
+from boto3.dynamodb.types import TypeDeserializer, Decimal
 
 from statikk.conditions import Condition, Equals, BeginsWith
 from statikk.expressions import UpdateExpressionBuilder
 from statikk.models import (
     DatabaseModel,
     GSI,
+    Index,
     IndexPrimaryKeyField,
     IndexSecondaryKeyField,
     KeySchema,
@@ -156,7 +157,13 @@ class Table:
         key = {self.key_schema.hash_key: id}
         self._get_dynamodb_table().delete_item(Key=key)
 
-    def get_item(self, id: str, model_class: Type[DatabaseModel], sort_key: Optional[Any] = None):
+    def get_item(
+        self,
+        id: str,
+        model_class: Type[DatabaseModel],
+        sort_key: Optional[Any] = None,
+        consistent_read: bool = False,
+    ):
         """
         Returns an item from the database by id, using the partition key of the table.
         :param id: The id of the item to retrieve.
@@ -166,11 +173,10 @@ class Table:
         key = {self.key_schema.hash_key: id}
         if sort_key:
             key[self.key_schema.sort_key] = self._serialize_value(sort_key)
-        raw_data = self._get_dynamodb_table().get_item(Key=key)
+        raw_data = self._get_dynamodb_table().get_item(Key=key, ConsistentRead=consistent_read)
         if "Item" not in raw_data:
             raise ItemNotFoundError(f"{model_class} with id '{id}' not found.")
         data = raw_data["Item"]
-        del data["type"]
         for key, value in data.items():
             data[key] = self._deserialize_value(value, model_class.model_fields[key])
         return model_class(**data)
@@ -206,7 +212,6 @@ class Table:
         """
         data = self._serialize_item(model)
         self._get_dynamodb_table().put_item(Item=data)
-        del data["type"]
         for key, value in data.items():
             data[key] = self._deserialize_value(value, model.model_fields[key])
         return type(model)(**data)
@@ -215,6 +220,7 @@ class Table:
         self,
         hash_key: str,
         update_builder: UpdateExpressionBuilder,
+        model: DatabaseModel,
         range_key: Optional[str] = None,
     ):
         (
@@ -226,16 +232,40 @@ class Table:
         if range_key:
             key[self.key_schema.sort_key] = range_key
 
+        def _find_changed_indexes():
+            changed_index_values = set()
+            for prefixed_attribute, value in expression_attribute_values.items():
+                expression_attribute_values[prefixed_attribute] = self._serialize_value(value)
+                attribute = prefixed_attribute.replace(":", "")
+                if model.model_fields[attribute].annotation is IndexSecondaryKeyField:
+                    idx_field = getattr(model, attribute)
+                    for idx in idx_field.index_names:
+                        idx_field.value = value
+                        changed_index_values.add(idx)
+            return changed_index_values
+
+        changed_index_values = _find_changed_indexes()
+
+        for idx_name in changed_index_values:
+            idx = [idx for idx in self.indexes if idx.name == idx_name][0]
+            index_value = self._get_sort_key_value(model, idx)
+            expression_attribute_values[f":{idx.sort_key.name}"] = index_value
+            update_expression += f" SET {idx.sort_key.name} = :{idx.sort_key.name}"
+
         request = {
             "Key": key,
             "UpdateExpression": update_expression,
             "ExpressionAttributeValues": expression_attribute_values,
+            "ReturnValues": "ALL_NEW",
         }
         if expression_attribute_names:
             request["ExpressionAttributeNames"] = expression_attribute_names
 
         response = self._get_dynamodb_table().update_item(**request)
-        return response
+        data = response["Attributes"]
+        for key, value in data.items():
+            data[key] = self._deserialize_value(value, model.model_fields[key])
+        return type(model)(**data)
 
     def batch_write(self):
         """
@@ -272,11 +302,11 @@ class Table:
             raise InvalidIndexNameError(f"The provided index name '{index_name}' is not configured on the table.")
         index = index_filter[0]
         key_condition = hash_key.evaluate(index.hash_key.name)
-        if range_key is None:
+        if range_key is None and model_class.include_type_in_sort_key():
             range_key = BeginsWith(model_class.model_type())
-
-        range_key.enrich(model_class=model_class)
-        key_condition = key_condition & range_key.evaluate(index.sort_key.name)
+        if range_key:
+            range_key.enrich(model_class=model_class)
+            key_condition = key_condition & range_key.evaluate(index.sort_key.name)
 
         query_params = {
             "IndexName": index_name,
@@ -288,13 +318,14 @@ class Table:
 
         while last_evaluated_key:
             items = self._get_dynamodb_table().query(**query_params)
-            yield from [model_class(**item) for item in items["Items"]]
+            yield from [self._deserialize_item(item, model_class=model_class) for item in items["Items"]]
             last_evaluated_key = items.get("LastEvaluatedKey", False)
 
     def scan(
         self,
         model_class: Type[DatabaseModel],
         filter_condition: Optional[ComparisonCondition] = None,
+        consistent_read: bool = False,
     ):
         """
         Scans the database for items matching the provided filter condition. The method returns a list of items matching
@@ -303,7 +334,9 @@ class Table:
         :param model_class: The model class to use to deserialize the items.
         :param filter_condition: An optional filter condition to use for the query. See boto3.dynamodb.conditions.ComparisonCondition for more information.
         """
-        query_params = {}
+        query_params = {
+            "ConsistentRead": consistent_read,
+        }
         if filter_condition:
             query_params["FilterExpression"] = filter_condition
         last_evaluated_key = True
@@ -349,11 +382,16 @@ class Table:
 
         return results
 
-    def _prepare_model_data(self, item: DatabaseModel) -> Dict[str, Any]:
-        for idx in self.indexes:
+    def _prepare_model_data(
+        self,
+        item: DatabaseModel,
+        indexes: List[Index],
+        force_override_index_fields: bool = False,
+    ) -> Dict[str, Any]:
+        for idx in indexes:
             index_fields = self._compose_index_values(item, idx)
             for key, value in index_fields.items():
-                if hasattr(item, key) and getattr(item, key) is not None:
+                if hasattr(item, key) and (getattr(item, key) is not None and not force_override_index_fields):
                     continue
                 if value is not None:
                     setattr(item, key, value)
@@ -361,20 +399,32 @@ class Table:
         return item.model_dump()
 
     def _serialize_item(self, item: DatabaseModel):
-        data = self._prepare_model_data(item)
+        data = self._prepare_model_data(item, self.indexes)
         for key, value in data.items():
             data[key] = self._serialize_value(value)
-        data["type"] = item.model_type()
         return data
 
+    def _deserialize_item(self, item: Dict[str, Any], model_class: Type[DatabaseModel]):
+        for key, value in item.items():
+            item[key] = self._deserialize_value(value, model_class.model_fields[key])
+        return model_class(**item)
+
     def _deserialize_value(self, value: Any, annotation: Any):
-        if annotation is datetime or "annotation=datetime" in str(annotation):
+        if annotation is datetime or "datetime" in str(annotation):
             return datetime.fromtimestamp(int(value))
+        if annotation is float:
+            return float(value)
         return value
 
     def _serialize_value(self, value: Any):
         if isinstance(value, datetime):
             return int(value.timestamp())
+        if isinstance(value, float):
+            return Decimal(value)
+        if isinstance(value, list):
+            return [self._serialize_value(item) for item in value]
+        if isinstance(value, dict):
+            return {key: self._serialize_value(item) for key, item in value.items() if item is not None}
         return value
 
     def _set_index_fields(self, model: DatabaseModel | Type[DatabaseModel], idx: GSI):
@@ -384,52 +434,55 @@ class Table:
         if idx.sort_key.name not in model_fields:
             model_fields[idx.sort_key.name] = FieldInfo(annotation=idx.sort_key.type, default=None, required=False)
 
-    def _compose_index_values(self, model: DatabaseModel, idx: GSI) -> Dict[str, Any]:
-        model_fields = model.model_fields
-        hash_key_field = [
-            field_name
-            for field_name, field_info in model_fields.items()
-            if field_info.annotation is not None
-            if field_info.annotation is IndexPrimaryKeyField and idx.name in getattr(model, field_name).index_names
-        ]
-        if len(hash_key_field) == 0 and model.type_is_primary_key():
-            hash_key_field.append(model.model_type())
-        hash_key_field = hash_key_field[0]
+    def _get_sort_key_value(self, model: DatabaseModel, idx: GSI) -> str:
         sort_key_fields_unordered = [
             (field_name, getattr(model, field_name).order)
-            for field_name, field_info in model_fields.items()
+            for field_name, field_info in model.model_fields.items()
             if field_info.annotation is not None
             if field_info.annotation is IndexSecondaryKeyField and idx.name in getattr(model, field_name).index_names
         ]
 
+        if len(sort_key_fields_unordered) == idx.sort_key is not None:
+            raise IncorrectSortKeyError(f"Model {model.__class__} does not have a sort key defined.")
         if sort_key_fields_unordered[0][1] is not None:
             sort_key_fields_unordered.sort(key=lambda x: x[1])
 
         sort_key_fields = [field[0] for field in sort_key_fields_unordered]
 
-        def _get_sort_key_value():
-            if len(sort_key_fields) == 0 and model.include_type_in_sort_key():
-                return model.model_type()
-            if idx.sort_key.type is not str:
-                value = getattr(model, sort_key_fields[0]).value
-                if type(value) is not idx.sort_key.type:
-                    raise IncorrectSortKeyError(
-                        f"Incorrect sort key type. Sort key type for sort key '{idx.sort_key.name}' should be: "
-                        + str(idx.sort_key.type)
-                        + " but got: "
-                        + str(type(value))
-                    )
-                value = value or idx.sort_key.default
-                return self._serialize_value(value)
+        if len(sort_key_fields) == 0 and model.include_type_in_sort_key():
+            return model.model_type()
+        if idx.sort_key.type is not str:
+            value = getattr(model, sort_key_fields[0]).value
+            if type(value) is not idx.sort_key.type:
+                raise IncorrectSortKeyError(
+                    f"Incorrect sort key type. Sort key type for sort key '{idx.sort_key.name}' should be: "
+                    + str(idx.sort_key.type)
+                    + " but got: "
+                    + str(type(value))
+                )
+            value = value or idx.sort_key.default
+            return self._serialize_value(value)
+
+        sort_key_values: List[str] = []
+        if model.include_type_in_sort_key() and model.model_type() not in sort_key_values:
+            sort_key_values.append(model.model_type())
+
+        for field in sort_key_fields:
+            value = getattr(model, field).value
+            sort_key_values.append(value)
+        return self.delimiter.join(sort_key_values)
 
-            sort_key_values: List[str] = []
-            if model.include_type_in_sort_key() and model.model_type() not in sort_key_values:
-                sort_key_values.append(model.model_type())
-
-            for field in sort_key_fields:
-                value = getattr(model, field).value
-                sort_key_values.append(value)
-            return self.delimiter.join(sort_key_values)
+    def _compose_index_values(self, model: DatabaseModel, idx: GSI) -> Dict[str, Any]:
+        model_fields = model.model_fields
+        hash_key_field = [
+            field_name
+            for field_name, field_info in model_fields.items()
+            if field_info.annotation is not None
+            if field_info.annotation is IndexPrimaryKeyField and idx.name in getattr(model, field_name).index_names
+        ]
+        if len(hash_key_field) == 0 and model.type_is_primary_key():
+            hash_key_field.append(model.model_type())
+        hash_key_field = hash_key_field[0]
 
         def _get_hash_key_value():
             if hash_key_field == model.model_type():
@@ -439,7 +492,7 @@ class Table:
 
         return {
             idx.hash_key.name: _get_hash_key_value(),
-            idx.sort_key.name: _get_sort_key_value(),
+            idx.sort_key.name: self._get_sort_key_value(model, idx),
         }
 
     def _perform_batch_write(self, put_items: List[DatabaseModel], delete_items: List[DatabaseModel]):

{statikk-0.0.8 → statikk-0.0.10}/src/statikk/expressions.py

@@ -591,10 +591,14 @@ class UpdateExpressionBuilder:
         self.model_schema = model_cls.model_json_schema()
         self.expression_attribute_names = {}
 
+    @classmethod
+    def safe_name(cls, key):
+        return f"#n_{key}"
+
     def _safe_name(self, key):
         """Replace reserved words with a safe placeholder."""
         if key.upper() in RESERVED_WORDS:
-            safe_key = f"#n_{key}"
+            safe_key = self.safe_name(key)
             self.expression_attribute_names[safe_key] = key
             return safe_key
         return key
@@ -669,11 +673,17 @@ class UpdateExpressionBuilder:
 
 
 class DatabaseModelUpdateExpressionBuilder(UpdateExpressionBuilder):
-    def __init__(self, model_cls: Type[DatabaseModel], id: str, sort_key: Optional[str] = None):
-        self.hash_key = id
+    def __init__(self, model, sort_key: Optional[str] = None):
+        self.model = model
+        self.hash_key = model.id
         self.sort_key = sort_key
-        self.model_cls = model_cls
-        super().__init__(model_cls)
+        self.model_cls = model.__class__
+        super().__init__(self.model_cls)
 
     def execute(self):
-        return self.model_cls.update_item(self.hash_key, range_key=self.sort_key, update_builder=self)
+        return self.model_cls.update_item(
+            self.hash_key,
+            range_key=self.sort_key,
+            update_builder=self,
+            model=self.model,
+        )

{statikk-0.0.8 → statikk-0.0.10}/src/statikk/models.py

@@ -70,9 +70,6 @@ class DatabaseModel(BaseModel):
     def batch_write(cls):
         return cls._table.batch_write()
 
-    class Config:
-        include_type_field_in_sort_key = False
-
     @classmethod
     def query(
         cls,
@@ -95,29 +92,38 @@ class DatabaseModel(BaseModel):
     def delete(self):
         return self._table.delete_item(self.id)
 
-    @classmethod
-    def update(cls, id: str, sort_key: Optional[str] = None) -> DatabaseModelUpdateExpressionBuilder:
-        return DatabaseModelUpdateExpressionBuilder(cls, id, sort_key)
+    def update(self, sort_key: Optional[str] = None) -> DatabaseModelUpdateExpressionBuilder:
+        return DatabaseModelUpdateExpressionBuilder(self, sort_key)
 
     @classmethod
     def update_item(
         cls,
         hash_key: str,
         update_builder: DatabaseModelUpdateExpressionBuilder,
+        model: DatabaseModel,
         range_key: Optional[str] = None,
     ):
-        return cls._table.update_item(hash_key, range_key=range_key, update_builder=update_builder)
+        return cls._table.update_item(
+            hash_key,
+            range_key=range_key,
+            update_builder=update_builder,
+            model=model,
+        )
 
     @classmethod
-    def get(cls, id: str, sort_key: Optional[str] = None):
-        return cls._table.get_item(id=id, model_class=cls, sort_key=sort_key)
+    def get(cls, id: str, sort_key: Optional[str] = None, consistent_read: bool = False):
+        return cls._table.get_item(id=id, model_class=cls, sort_key=sort_key, consistent_read=consistent_read)
 
     @classmethod
     def batch_get(cls, ids: List[str], batch_size: int = 100):
         return cls._table.batch_get_items(ids=ids, model_class=cls, batch_size=batch_size)
 
     @classmethod
-    def scan(cls, filter_condition: Optional[ComparisonCondition] = None):
+    def scan(
+        cls,
+        filter_condition: Optional[ComparisonCondition] = None,
+        consistent_read: bool = False,
+    ):
         return cls._table.scan(model_class=cls, filter_condition=filter_condition)
 
     @staticmethod

{statikk-0.0.8 → statikk-0.0.10}/src/statikk.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: statikk
-Version: 0.0.8
+Version: 0.0.10
 Summary: statikk is a single table application (STA) library for DynamoDb.
 Home-page: https://github.com/terinia/statikk
 Author: Balint Biro

{statikk-0.0.8 → statikk-0.0.10}/tests/test_engine.py

@@ -1,6 +1,7 @@
+from _decimal import Decimal
 from datetime import datetime, timezone
 from typing import List
-
+from pydantic import BaseModel
 import pytest
 from boto3.dynamodb.conditions import Attr
 from moto import mock_dynamodb
@@ -28,6 +29,7 @@ class MyAwesomeModel(DatabaseModel):
     name: str = "Foo"
     values: set = {1, 2, 3, 4}
     cost: int = 4
+    probability: float = 0.5
 
 
 class SimpleModel(DatabaseModel):
@@ -89,6 +91,7 @@ def test_create_my_awesome_model():
         "values": {1, 2, 3, 4},
         "cost": 4,
         "type": "MyAwesomeModel",
+        "probability": 0.5,
     }
     model_2 = MyAwesomeModel(id="foo-2", player_id="123", tier="EPIC", name="FooFoo")
     table.put_item(model_2)
@@ -102,6 +105,7 @@ def test_create_my_awesome_model():
         "values": {1, 2, 3, 4},
         "cost": 4,
         "type": "MyAwesomeModel",
+        "probability": 0.5,
     }
     mock_dynamodb().stop()
 
@@ -591,11 +595,7 @@ def test_get_item_does_not_exist():
     mock_dynamodb().stop()
 
 
-def test_update_set_attribute():
-    pass
-
-
-def test_update_add_attribute():
+def test_update():
     mock_dynamodb().start()
     table = Table(
         name="my-dynamodb-table",
@@ -612,13 +612,23 @@ def test_update_add_attribute():
     _create_dynamodb_table(table)
     model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY", name="FooFoo", values={1, 2, 3, 4})
     model.save()
-    MyAwesomeModel.update("foo").set("player_id", "456").delete("values", {1}).remove("name").add("cost", 1).execute()
+    (
+        model.update()
+        .set("player_id", "456")
+        .set("tier", "EPIC")
+        .delete("values", {1})
+        .remove("name")
+        .add("cost", 1)
+        .execute()
+    )
     item = table.get_item("foo", MyAwesomeModel)
     assert item.player_id.value == "456"
     assert item.values == {2, 3, 4}
     assert item.name == "Foo"  # default value
     assert item.cost == 5
-    MyAwesomeModel.update("foo").set("name", "FooFoo").execute()
+    assert item.tier.value == "EPIC"
+    assert item.gsi_sk == "MyAwesomeModel|EPIC"
+    item.update().set("name", "FooFoo").execute()
     item = table.get_item("foo", MyAwesomeModel)
     assert item.name == "FooFoo"
     mock_dynamodb().stop()
@@ -676,6 +686,41 @@ def test_query_no_range_key_provided():
     mock_dynamodb().stop()
 
 
+def test_query_no_range_is_provided_but_model_does_not_include_type_in_range_key():
+    class Model(DatabaseModel):
+        tier: IndexSecondaryKeyField
+
+        @classmethod
+        def type_is_primary_key(cls):
+            return True
+
+        @classmethod
+        def include_type_in_sort_key(cls):
+            return False
+
+    mock_dynamodb().start()
+    table = Table(
+        name="my-dynamodb-table",
+        key_schema=KeySchema(hash_key="id"),
+        indexes=[
+            GSI(
+                name="main-index",
+                hash_key=Key(name="gsi_pk"),
+                sort_key=Key(name="gsi_sk"),
+            )
+        ],
+        models=[Model],
+    )
+    _create_dynamodb_table(table)
+    m = Model(tier="LEGENDARY")
+    m.save()
+    models = list(Model.query(hash_key=Equals("Model")))
+    assert len(models) == 1
+    assert models[0].tier.value == "LEGENDARY"
+    assert models[0].gsi_pk == "Model"
+    assert models[0].gsi_sk == "LEGENDARY"
+
+
 def test_index_field_order_is_respected():
     class ModelWithIndexOrdersDefined(DatabaseModel):
         player_id: IndexPrimaryKeyField
@@ -703,3 +748,71 @@ def test_index_field_order_is_respected():
     model.save()
     item = table.get_item("123", ModelWithIndexOrdersDefined)
     assert item.gsi_sk == "ModelWithIndexOrdersDefined|EPIC|Mage"
+
+
+def test_nested_models():
+    class InnerInnerModel(BaseModel):
+        baz: str
+
+    class InnerModel(BaseModel):
+        foo: str
+        values: List[datetime] = [
+            datetime(2023, 9, 9, 12, 0, 0),
+            datetime(2023, 9, 9, 13, 0, 0),
+        ]
+        cost: int = 5
+        inner_inner: InnerInnerModel
+
+    class NestedModel(DatabaseModel):
+        player_id: IndexPrimaryKeyField
+        unit_class: IndexSecondaryKeyField = IndexSecondaryKeyField(order=2)
+        tier: IndexSecondaryKeyField = IndexSecondaryKeyField(order=1)
+        name: str = "Foo"
+        values: set = {1, 2, 3, 4}
+        cost: int = 4
+        inner_model: InnerModel
+
+    mock_dynamodb().start()
+    table = Table(
+        name="my-dynamodb-table",
+        key_schema=KeySchema(hash_key="id"),
+        indexes=[
+            GSI(
+                name="main-index",
+                hash_key=Key(name="gsi_pk"),
+                sort_key=Key(name="gsi_sk"),
+            )
+        ],
+        models=[NestedModel],
+    )
+    _create_dynamodb_table(table)
+    model = NestedModel(
+        id="123",
+        player_id="456",
+        unit_class="Mage",
+        tier="EPIC",
+        inner_model=InnerModel(foo="bar", inner_inner=InnerInnerModel(baz="baz")),
+    )
+    model.save()
+    item = table.get_item("123", NestedModel)
+    assert item.model_dump() == {
+        "id": "123",
+        "player_id": "456",
+        "unit_class": "Mage",
+        "tier": "EPIC",
+        "name": "Foo",
+        "values": {Decimal("1"), Decimal("2"), Decimal("3"), Decimal("4")},
+        "cost": 4,
+        "inner_model": {
+            "foo": "bar",
+            "values": [
+                datetime(2023, 9, 9, 10, 0, tzinfo=timezone.utc),
+                datetime(2023, 9, 9, 11, 0, tzinfo=timezone.utc),
+            ],
+            "cost": 5,
+            "inner_inner": {"baz": "baz"},
+        },
+        "gsi_pk": "456",
+        "gsi_sk": "NestedModel|EPIC|Mage",
+        "type": "NestedModel",
+    }