# This file is execfile()d with the current directory set to its containing dir.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html
#
# All configuration values have a default; values that are commented out
# serve to show the default.

import os
import sys
import shutil

# -- Path setup --------------------------------------------------------------

__location__ = os.path.dirname(__file__)

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
sys.path.insert(0, os.path.join(__location__, "../src"))

# -- Run sphinx-apidoc -------------------------------------------------------
# This hack is necessary since RTD does not issue `sphinx-apidoc` before running
# `sphinx-build -b html . _build/html`. See Issue:
# https://github.com/readthedocs/readthedocs.org/issues/1139
# DON'T FORGET: Check the box "Install your project inside a virtualenv using
# setup.py install" in the RTD Advanced Settings.
# Additionally it helps us to avoid running apidoc manually

try:  # for Sphinx >= 1.7
    from sphinx.ext import apidoc
except ImportError:
    from sphinx import apidoc

output_dir = os.path.join(__location__, "api")
module_dir = os.path.join(__location__, "../src/statikk")
try:
    shutil.rmtree(output_dir)
except FileNotFoundError:
    pass

try:
    import sphinx

    cmd_line = f"sphinx-apidoc --implicit-namespaces -f -o {output_dir} {module_dir}"

    args = cmd_line.split(" ")
    if tuple(sphinx.__version__.split(".")) >= ("1", "7"):
        # This is a rudimentary parse_version to avoid external dependencies
        args = args[1:]

    apidoc.main(args)
except Exception as e:
    print("Running `sphinx-apidoc` failed!\n{}".format(e))

# -- General configuration ---------------------------------------------------

# If your documentation needs a minimal Sphinx version, state it here.
# needs_sphinx = '1.0'

# Add any Sphinx extension module names here, as strings. They can be extensions
# coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
extensions = [
    "sphinx.ext.autodoc",
    "sphinx.ext.intersphinx",
    "sphinx.ext.todo",
    "sphinx.ext.autosummary",
    "sphinx.ext.viewcode",
    "sphinx.ext.coverage",
    "sphinx.ext.doctest",
    "sphinx.ext.ifconfig",
    "sphinx.ext.mathjax",
    "sphinx.ext.napoleon",
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# The suffix of source filenames.
source_suffix = ".rst"

# The encoding of source files.
# source_encoding = 'utf-8-sig'

# The master toctree document.
master_doc = "index"

# General information about the project.
project = "statikk"
copyright = "2023, Balint Biro"

# The version info for the project you're documenting, acts as replacement for
# |version| and |release|, also used in various other places throughout the
# built documents.
#
# version: The short X.Y version.
# release: The full version, including alpha/beta/rc tags.
# If you don’t need the separation provided between version and release,
# just set them both to the same value.
try:
    from statikk import __version__ as version
except ImportError:
    version = ""

if not version or version.lower() == "unknown":
    version = os.getenv("READTHEDOCS_VERSION", "unknown")  # automatically set by RTD

release = version

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
# language = None

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
# today = ''
# Else, today_fmt is used as the format for a strftime call.
# today_fmt = '%B %d, %Y'

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store", ".venv"]

# The reST default role (used for this markup: `text`) to use for all documents.
# default_role = None

# If true, '()' will be appended to :func: etc. cross-reference text.
# add_function_parentheses = True

# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
# add_module_names = True

# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
# show_authors = False

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = "sphinx"

# A list of ignored prefixes for module index sorting.
# modindex_common_prefix = []

# If true, keep warnings as "system message" paragraphs in the built documents.
# keep_warnings = False

# If this is True, todo emits a warning for each TODO entries. The default is False.
todo_emit_warnings = True


# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
html_theme = "alabaster"

# Theme options are theme-specific and customize the look and feel of a theme
# further.  For a list of options available for each theme, see the
# documentation.
html_theme_options = {"sidebar_width": "300px", "page_width": "1200px"}

# Add any paths that contain custom themes here, relative to this directory.
# html_theme_path = []

# The name for this set of Sphinx documents.  If None, it defaults to
# "<project> v<release> documentation".
# html_title = None

# A shorter title for the navigation bar.  Default is the same as html_title.
# html_short_title = None

# The name of an image file (relative to this directory) to place at the top
# of the sidebar.
# html_logo = ""

# The name of an image file (within the static path) to use as favicon of the
# docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
# pixels large.
# html_favicon = None

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]

# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
# using the given strftime format.
# html_last_updated_fmt = '%b %d, %Y'

# If true, SmartyPants will be used to convert quotes and dashes to
# typographically correct entities.
# html_use_smartypants = True

# Custom sidebar templates, maps document names to template names.
# html_sidebars = {}

# Additional templates that should be rendered to pages, maps page names to
# template names.
# html_additional_pages = {}

# If false, no module index is generated.
# html_domain_indices = True

# If false, no index is generated.
# html_use_index = True

# If true, the index is split into individual pages for each letter.
# html_split_index = False

# If true, links to the reST sources are added to the pages.
# html_show_sourcelink = True

# If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
# html_show_sphinx = True

# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
# html_show_copyright = True

# If true, an OpenSearch description file will be output, and all pages will
# contain a <link> tag referring to it.  The value of this option must be the
# base URL from which the finished HTML is served.
# html_use_opensearch = ''

# This is the file name suffix for HTML files (e.g. ".xhtml").
# html_file_suffix = None

# Output file base name for HTML help builder.
htmlhelp_basename = "statikk-doc"


# -- Options for LaTeX output ------------------------------------------------

latex_elements = {
    # The paper size ("letterpaper" or "a4paper").
    # "papersize": "letterpaper",
    # The font size ("10pt", "11pt" or "12pt").
    # "pointsize": "10pt",
    # Additional stuff for the LaTeX preamble.
    # "preamble": "",
}

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title, author, documentclass [howto/manual]).
latex_documents = [("index", "user_guide.tex", "statikk Documentation", "Balint Biro", "manual")]

# The name of an image file (relative to this directory) to place at the top of
# the title page.
# latex_logo = ""

# For "manual" documents, if this is true, then toplevel headings are parts,
# not chapters.
# latex_use_parts = False

# If true, show page references after internal links.
# latex_show_pagerefs = False

# If true, show URL addresses after external links.
# latex_show_urls = False

# Documents to append as an appendix to all manuals.
# latex_appendices = []

# If false, no module index is generated.
# latex_domain_indices = True

# -- External mapping --------------------------------------------------------
python_version = ".".join(map(str, sys.version_info[0:2]))
intersphinx_mapping = {
    "sphinx": ("https://www.sphinx-doc.org/en/master", None),
    "python": ("https://docs.python.org/" + python_version, None),
    "matplotlib": ("https://matplotlib.org", None),
    "numpy": ("https://numpy.org/doc/stable", None),
    "sklearn": ("https://scikit-learn.org/stable", None),
    "pandas": ("https://pandas.pydata.org/pandas-docs/stable", None),
    "scipy": ("https://docs.scipy.org/doc/scipy/reference", None),
    "setuptools": ("https://setuptools.pypa.io/en/stable/", None),
    "pyscaffold": ("https://pyscaffold.org/en/stable", None),
}

print(f"loading configurations for {project} {version} ...", file=sys.stderr)

docs/conf.py
terinia/statikk
"""
    Setup file for statikk.
    Use setup.cfg to configure your project.

    This file was generated with PyScaffold 4.5.
    PyScaffold helps you to put up the scaffold of your new Python project.
    Learn more under: https://pyscaffold.org/
"""
from setuptools import setup

if __name__ == "__main__":
    try:
        setup(use_scm_version={"version_scheme": "no-guess-dev"})
    except:  # noqa
        print(
            "\n\nAn error occurred while building the project, "
            "please ensure you have the most updated version of setuptools, "
            "setuptools_scm and wheel with:\n"
            "   pip install -U setuptools setuptools_scm wheel\n\n"
        )
        raise

setup.py
terinia/statikk
"""
This module contains the Condition class and its subclasses. These classes are used to construct expressions understood
by DynamoDB to narrow down query results. Not to be confused with the ConditionExpression parameter, which is used to
narrow down filter results and is not implemented in this library. This library relies on the boto3 library to handle
the ConditionExpression filter_condition arguments.
See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/customizations/dynamodb.html#boto3.dynamodb.conditions.Attr
for more information on filter expressions.

Example usage of this module:
    from statikk.conditions import Equals, BeginsWith
    app.query(range_key=Equals("123"), hash_key=BeginsWith("abc"))
"""
from abc import ABC, abstractmethod
from boto3.dynamodb.conditions import Key, ComparisonCondition
from typing import Any


class Condition(ABC):
    def __init__(self, value: Any):
        self.value = value

    @abstractmethod
    def evaluate(self, key: Any) -> ComparisonCondition:
        pass

    def enrich(self, **kwargs):
        pass


class Equals(Condition):
    def evaluate(self, key: Any) -> ComparisonCondition:
        return Key(key).eq(self.value)


class BeginsWith(Condition):
    def evaluate(self, key: Any) -> ComparisonCondition:
        return Key(key).begins_with(self.value)

    def enrich(self, model_class, **kwargs):
        if not self.value.startswith(model_class.model_type()) and model_class.include_type_in_sort_key():
            self.value = f"{model_class.model_type()}|{self.value}"


class LessThan(Condition):
    def evaluate(self, key: Any) -> ComparisonCondition:
        return Key(key).lt(self.value)


class GreaterThan(Condition):
    def evaluate(self, key: Any) -> ComparisonCondition:
        return Key(key).gt(self.value)


class LessThanOrEqual(Condition):
    def evaluate(self, key: Any) -> ComparisonCondition:
        return Key(key).lte(self.value)


class GreaterThanOrEqual(Condition):
    def evaluate(self, key: Any) -> ComparisonCondition:
        return Key(key).gte(self.value)


class Between(Condition):
    def evaluate(self, key: Any) -> ComparisonCondition:
        return Key(key).between(*self.value)

src/statikk/conditions.py
terinia/statikk
import os
from datetime import datetime
from typing import Any, Dict, Type, Optional, List, Union

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, Decimal

from statikk.conditions import Condition, Equals, BeginsWith
from statikk.expressions import UpdateExpressionBuilder
from statikk.models import (
    DatabaseModel,
    GSI,
    Index,
    IndexPrimaryKeyField,
    IndexSecondaryKeyField,
    KeySchema,
)


class InvalidIndexNameError(Exception):
    pass


class IncorrectSortKeyError(Exception):
    pass


class ItemNotFoundError(Exception):
    pass


class Table:
    def __init__(
        self,
        name: str,
        models: List[Type[DatabaseModel]],
        key_schema: KeySchema,
        indexes: List[GSI] = Optional[None],
        delimiter: str = "|",
        billing_mode: str = "PAY_PER_REQUEST",
    ):
        self.name = name
        self.key_schema = key_schema
        self.indexes = indexes or []
        self.delimiter = delimiter
        self.models = models
        self.billing_mode = billing_mode
        for idx in self.indexes:
            for model in self.models:
                self._set_index_fields(model, idx)
                model.set_table_ref(self)
                if "type" not in model.model_fields:
                    model.model_fields["type"] = FieldInfo(annotation=str, default=model.model_type(), required=False)

    def _dynamodb_client(self):
        return boto3.client(
            "dynamodb",
            config=Config(region_name=os.environ.get("AWS_DEFAULT_REGION", "eu-west-1")),
        )

    def _get_dynamodb_table(self):
        dynamodb = boto3.resource(
            "dynamodb",
            config=Config(region_name=os.environ.get("AWS_DEFAULT_REGION", "eu-west-1")),
        )
        return dynamodb.Table(self.name)

    def _to_dynamodb_type(self, type: Any):
        if type is str:
            return "S"
        if type is int:
            return "N"
        if type is datetime:
            return "N"
        if type is float:
            return "N"
        if type is bool:
            return "BOOL"
        if type is list:
            return "L"
        if type is dict:
            return "M"
        if type is None:
            return "NULL"
        raise ValueError(f"Unsupported type: {type}.")

    def create(self, aws_region: Optional[str] = None):
        """Creates a DynamoDB table from the specified definition.

        This method only supports a subset of all available configuration values on DynamoDB tables.
        It is recommended that you manage your DynamoDb table outside the scope of your application, using CloudFormation or Terraform.

        This method is mostly here to provide a shorthand for boostrapping an in-memory DynamoDb table mock
        in moto-based unit and integration tests.
        """
        region = aws_region or os.environ.get("AWS_DEFAULT_REGION")
        if not region:
            raise ValueError(
                "AWS region not specified. Please provide a region or set the AWS_DEFAULT_REGION environment variable."
            )
        hash_key_attribute_definitions = [
            {
                "AttributeName": index.hash_key.name,
                "AttributeType": self._to_dynamodb_type(index.hash_key.type),
            }
            for index in self.indexes
        ]
        sort_key_attribute_definitions = [
            {
                "AttributeName": index.sort_key.name,
                "AttributeType": self._to_dynamodb_type(index.sort_key.type),
            }
            for index in self.indexes
        ]
        global_secondary_indexes = [
            {
                "IndexName": index.name,
                "KeySchema": [
                    {
                        "AttributeName": index.hash_key.name,
                        "KeyType": "HASH",
                    },
                    {
                        "AttributeName": index.sort_key.name,
                        "KeyType": "RANGE",
                    },
                ],
                "Projection": {"ProjectionType": index.projection_type},
            }
            for index in self.indexes
        ]
        key_schema = [{"AttributeName": self.key_schema.hash_key, "KeyType": "HASH"}]
        if self.key_schema.sort_key:
            key_schema.append({"AttributeName": self.key_schema.sort_key, "KeyType": "RANGE"})
        self._dynamodb_client().create_table(
            TableName=self.name,
            KeySchema=key_schema,
            AttributeDefinitions=[{"AttributeName": self.key_schema.hash_key, "AttributeType": "S"}]
            + hash_key_attribute_definitions
            + sort_key_attribute_definitions,
            GlobalSecondaryIndexes=global_secondary_indexes,
            BillingMode=self.billing_mode,
        )

    def delete(self):
        """Deletes the DynamoDB table."""
        self._dynamodb_client().delete_table(TableName=self.name)

    def delete_item(self, id: str):
        """
        Deletes an item from the database by id, using the partition key of the table.
        :param id: The id of the item to delete.
        """
        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,
        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.
        :param model_class: The model class to use to deserialize the item.
        :param sort_key: The sort key of the item to retrieve. If the table does not have a sort key, this parameter should not be provided.
        """
        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, ConsistentRead=consistent_read)
        if "Item" not in raw_data:
            raise ItemNotFoundError(f"{model_class} with id '{id}' not found.")
        data = raw_data["Item"]
        for key, value in data.items():
            data[key] = self._deserialize_value(value, model_class.model_fields[key])
        return model_class(**data)

    def put_item(self, model: DatabaseModel) -> DatabaseModel:
        """
        Puts an item into the database.

        This method automatically constructs values for the index fields of the provided model_class, based on the configured
        indexes on the table. If the model_class already has a value for the index field, the field's value is going to be used,
        otherwise the value is going to be constructed based on the model's fields. If a type field is not provided on the model class
        or the model instance, the type field is going to be set to the model class name. In case of string values, the values are
        always prefixed with the type of the model to avoid collisions between different model types.

        Example:
            class Card(DatabaseModel):
                id: str
                player_id: IndexPrimaryKeyField
                type: IndexSecondaryKeyField
                tier: IndexSecondaryKeyField

            table = Table(indexes=[GSI(hash_key=Key("gsi_pk"), sort_key=Key("gsi_sk"))])

            The constructed index fields for the Card model are going to be:
                gsi_pk: <player_id>
                gsi_sk: card|<tier>

            A setup like this allows for queries to select for all cards of a player with a specific tier, avoiding potential
            collisions with models where the hash_key is also the player_id.


        Returns the enriched database model instance.
        """
        data = self._serialize_item(model)
        self._get_dynamodb_table().put_item(Item=data)
        for key, value in data.items():
            data[key] = self._deserialize_value(value, model.model_fields[key])
        return type(model)(**data)

    def update_item(
        self,
        hash_key: str,
        update_builder: UpdateExpressionBuilder,
        model: DatabaseModel,
        range_key: Optional[str] = None,
    ):
        (
            update_expression,
            expression_attribute_values,
            expression_attribute_names,
        ) = update_builder.build()
        key = {self.key_schema.hash_key: hash_key}
        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)
        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):
        """
        Returns a context manager for batch writing items to the database. This method handles all the buffering of the
        batch operation and the construction of index fields for each item.
        """
        return BatchWriteContext(self)

    def query_index(
        self,
        hash_key: Union[Condition | str],
        model_class: Type[DatabaseModel],
        range_key: Optional[Condition] = None,
        filter_condition: Optional[ComparisonCondition] = None,
        index_name: Optional[str] = None,
    ):
        """
        Queries the database using the provided hash key and range key conditions. A filter condition can also be provided
        using the filter_condition parameter. The method returns a list of items matching the query, deserialized into the
        provided model_class parameter.

        :param hash_key: The hash key condition to use for the query. See statikk.conditions.Condition for more information.
        :param range_key: The range key condition to use for the query. See statikk.conditions.Condition for more information.
        :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.
        :param index_name: The name of the index to use for the query. If not provided, the first index configured on the table is used.
        """
        if isinstance(hash_key, str):
            hash_key = Equals(hash_key)
        if not index_name:
            index_name = self.indexes[0].name
        index_filter = [idx for idx in self.indexes if idx.name == index_name]
        if not index_filter:
            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 and model_class.include_type_in_sort_key():
            range_key = BeginsWith(model_class.model_type())
        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,
            "KeyConditionExpression": key_condition,
        }
        if filter_condition:
            query_params["FilterExpression"] = filter_condition
        last_evaluated_key = True

        while last_evaluated_key:
            items = self._get_dynamodb_table().query(**query_params)
            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
        the query, deserialized into the provided model_class parameter.

        :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 = {
            "ConsistentRead": consistent_read,
        }
        if filter_condition:
            query_params["FilterExpression"] = filter_condition
        last_evaluated_key = True

        while last_evaluated_key:
            items = self._get_dynamodb_table().scan(**query_params)
            yield from [model_class(**item) for item in items["Items"]]
            last_evaluated_key = items.get("LastEvaluatedKey", False)

    def _convert_dynamodb_to_python(self, item) -> Dict[str, Any]:
        deserializer = TypeDeserializer()
        return {k: deserializer.deserialize(v) for k, v in item.items()}

    def batch_get_items(
        self, ids: List[str], model_class: Type[DatabaseModel], batch_size: int = 100
    ) -> List[DatabaseModel]:
        dynamodb = self._dynamodb_client()

        id_batches = [ids[i : i + batch_size] for i in range(0, len(ids), batch_size)]

        results = []

        for batch in id_batches:
            request_items = {
                self.name: {
                    "Keys": [{"id": {"S": id}} for id in batch],
                }
            }

            while request_items:
                response = dynamodb.batch_get_item(RequestItems=request_items)

                if "UnprocessedKeys" in response and response["UnprocessedKeys"]:
                    request_items = response["UnprocessedKeys"][self.name]
                else:
                    results.extend(
                        [
                            model_class(**self._convert_dynamodb_to_python(item))
                            for item in response["Responses"][self.name]
                        ]
                    )
                    break

        return results

    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 and not force_override_index_fields):
                    continue
                if value is not None:
                    setattr(item, key, value)
        item.model_rebuild(force=True)
        return item.model_dump()

    def _serialize_item(self, item: DatabaseModel):
        data = self._prepare_model_data(item, self.indexes)
        for key, value in data.items():
            data[key] = self._serialize_value(value)
        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 "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):
        model_fields = model.model_fields
        if idx.hash_key.name not in model_fields:
            model_fields[idx.hash_key.name] = FieldInfo(annotation=idx.hash_key.type, default=None, required=False)
        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 _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.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]

        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)

    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():
                return model.model_type()
            else:
                return getattr(model, hash_key_field).value

        return {
            idx.hash_key.name: _get_hash_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]):
        if len(put_items) == 0 and len(delete_items) == 0:
            return

        dynamodb_table = self._get_dynamodb_table()

        if len(put_items) > 0:
            with dynamodb_table.batch_writer() as batch:
                for item in put_items:
                    data = self._serialize_item(item)
                    batch.put_item(Item=data)

        if len(delete_items) > 0:
            with dynamodb_table.batch_writer() as batch:
                for item in delete_items:
                    data = self._serialize_item(item)
                    batch.delete_item(Key=data)


class BatchWriteContext:
    def __init__(self, app: Table):
        self._table = app
        self._put_items: List[DatabaseModel] = []
        self._delete_items: List[DatabaseModel] = []

    def put(self, item: DatabaseModel):
        self._put_items.append(item)

    def delete(self, item: DatabaseModel):
        self._delete_items.append(item)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._table._perform_batch_write(self._put_items, self._delete_items)

src/statikk/engine.py
terinia/statikk
from __future__ import annotations
from typing import Type, TYPE_CHECKING, Optional, Tuple, Any

if TYPE_CHECKING:
    from statikk.models import DatabaseModel

RESERVED_WORDS = [
    "ABORT",
    "ABSOLUTE",
    "ACTION",
    "ADD",
    "AFTER",
    "AGENT",
    "AGGREGATE",
    "ALL",
    "ALLOCATE",
    "ALTER",
    "ANALYZE",
    "AND",
    "ANY",
    "ARCHIVE",
    "ARE",
    "ARRAY",
    "AS",
    "ASC",
    "ASCII",
    "ASENSITIVE",
    "ASSERTION",
    "ASYMMETRIC",
    "AT",
    "ATOMIC",
    "ATTACH",
    "ATTRIBUTE",
    "AUTH",
    "AUTHORIZATION",
    "AUTHORIZE",
    "AUTO",
    "AVG",
    "BACK",
    "BACKUP",
    "BASE",
    "BATCH",
    "BEFORE",
    "BEGIN",
    "BETWEEN",
    "BIGINT",
    "BINARY",
    "BIT",
    "BLOB",
    "BLOCK",
    "BOOLEAN",
    "BOTH",
    "BREADTH",
    "BUCKET",
    "BULK",
    "BY",
    "BYTE",
    "CALL",
    "CALLED",
    "CALLING",
    "CAPACITY",
    "CASCADE",
    "CASCADED",
    "CASE",
    "CAST",
    "CATALOG",
    "CHAR",
    "CHARACTER",
    "CHECK",
    "CLASS",
    "CLOB",
    "CLOSE",
    "CLUSTER",
    "CLUSTERED",
    "CLUSTERING",
    "CLUSTERS",
    "COALESCE",
    "COLLATE",
    "COLLATION",
    "COLLECTION",
    "COLUMN",
    "COLUMNS",
    "COMBINE",
    "COMMENT",
    "COMMIT",
    "COMPACT",
    "COMPILE",
    "COMPRESS",
    "CONDITION",
    "CONFLICT",
    "CONNECT",
    "CONNECTION",
    "CONSISTENCY",
    "CONSISTENT",
    "CONSTRAINT",
    "CONSTRAINTS",
    "CONSTRUCTOR",
    "CONSUMED",
    "CONTINUE",
    "CONVERT",
    "COPY",
    "CORRESPONDING",
    "COUNT",
    "COUNTER",
    "CREATE",
    "CROSS",
    "CUBE",
    "CURRENT",
    "CURSOR",
    "CYCLE",
    "DATA",
    "DATABASE",
    "DATE",
    "DATETIME",
    "DAY",
    "DEALLOCATE",
    "DEC",
    "DECIMAL",
    "DECLARE",
    "DEFAULT",
    "DEFERRABLE",
    "DEFERRED",
    "DEFINE",
    "DEFINED",
    "DEFINITION",
    "DELETE",
    "DELIMITED",
    "DEPTH",
    "DEREF",
    "DESC",
    "DESCRIBE",
    "DESCRIPTOR",
    "DETACH",
    "DETERMINISTIC",
    "DIAGNOSTICS",
    "DIRECTORIES",
    "DISABLE",
    "DISCONNECT",
    "DISTINCT",
    "DISTRIBUTE",
    "DO",
    "DOMAIN",
    "DOUBLE",
    "DROP",
    "DUMP",
    "DURATION",
    "DYNAMIC",
    "EACH",
    "ELEMENT",
    "ELSE",
    "ELSEIF",
    "EMPTY",
    "ENABLE",
    "END",
    "EQUAL",
    "EQUALS",
    "ERROR",
    "ESCAPE",
    "ESCAPED",
    "EVAL",
    "EVALUATE",
    "EXCEEDED",
    "EXCEPT",
    "EXCEPTION",
    "EXCEPTIONS",
    "EXCLUSIVE",
    "EXEC",
    "EXECUTE",
    "EXISTS",
    "EXIT",
    "EXPLAIN",
    "EXPLODE",
    "EXPORT",
    "EXPRESSION",
    "EXTENDED",
    "EXTERNAL",
    "EXTRACT",
    "FAIL",
    "FALSE",
    "FAMILY",
    "FETCH",
    "FIELDS",
    "FILE",
    "FILTER",
    "FILTERING",
    "FINAL",
    "FINISH",
    "FIRST",
    "FIXED",
    "FLATTERN",
    "FLOAT",
    "FOR",
    "FORCE",
    "FOREIGN",
    "FORMAT",
    "FORWARD",
    "FOUND",
    "FREE",
    "FROM",
    "FULL",
    "FUNCTION",
    "FUNCTIONS",
    "GENERAL",
    "GENERATE",
    "GET",
    "GLOB",
    "GLOBAL",
    "GO",
    "GOTO",
    "GRANT",
    "GREATER",
    "GROUP",
    "GROUPING",
    "HANDLER",
    "HASH",
    "HAVE",
    "HAVING",
    "HEAP",
    "HIDDEN",
    "HOLD",
    "HOUR",
    "IDENTIFIED",
    "IDENTITY",
    "IF",
    "IGNORE",
    "IMMEDIATE",
    "IMPORT",
    "IN",
    "INCLUDING",
    "INCLUSIVE",
    "INCREMENT",
    "INCREMENTAL",
    "INDEX",
    "INDEXED",
    "INDEXES",
    "INDICATOR",
    "INFINITE",
    "INITIALLY",
    "INLINE",
    "INNER",
    "INNTER",
    "INOUT",
    "INPUT",
    "INSENSITIVE",
    "INSERT",
    "INSTEAD",
    "INT",
    "INTEGER",
    "INTERSECT",
    "INTERVAL",
    "INTO",
    "INVALIDATE",
    "IS",
    "ISOLATION",
    "ITEM",
    "ITEMS",
    "ITERATE",
    "JOIN",
    "KEY",
    "KEYS",
    "LAG",
    "LANGUAGE",
    "LARGE",
    "LAST",
    "LATERAL",
    "LEAD",
    "LEADING",
    "LEAVE",
    "LEFT",
    "LENGTH",
    "LESS",
    "LEVEL",
    "LIKE",
    "LIMIT",
    "LIMITED",
    "LINES",
    "LIST",
    "LOAD",
    "LOCAL",
    "LOCALTIME",
    "LOCALTIMESTAMP",
    "LOCATION",
    "LOCATOR",
    "LOCK",
    "LOCKS",
    "LOG",
    "LOGED",
    "LONG",
    "LOOP",
    "LOWER",
    "MAP",
    "MATCH",
    "MATERIALIZED",
    "MAX",
    "MAXLEN",
    "MEMBER",
    "MERGE",
    "METHOD",
    "METRICS",
    "MIN",
    "MINUS",
    "MINUTE",
    "MISSING",
    "MOD",
    "MODE",
    "MODIFIES",
    "MODIFY",
    "MODULE",
    "MONTH",
    "MULTI",
    "MULTISET",
    "NAME",
    "NAMES",
    "NATIONAL",
    "NATURAL",
    "NCHAR",
    "NCLOB",
    "NEW",
    "NEXT",
    "NO",
    "NONE",
    "NOT",
    "NULL",
    "NULLIF",
    "NUMBER",
    "NUMERIC",
    "OBJECT",
    "OF",
    "OFFLINE",
    "OFFSET",
    "OLD",
    "ON",
    "ONLINE",
    "ONLY",
    "OPAQUE",
    "OPEN",
    "OPERATOR",
    "OPTION",
    "OR",
    "ORDER",
    "ORDINALITY",
    "OTHER",
    "OTHERS",
    "OUT",
    "OUTER",
    "OUTPUT",
    "OVER",
    "OVERLAPS",
    "OVERRIDE",
    "OWNER",
    "PAD",
    "PARALLEL",
    "PARAMETER",
    "PARAMETERS",
    "PARTIAL",
    "PARTITION",
    "PARTITIONED",
    "PARTITIONS",
    "PATH",
    "PERCENT",
    "PERCENTILE",
    "PERMISSION",
    "PERMISSIONS",
    "PIPE",
    "PIPELINED",
    "PLAN",
    "POOL",
    "POSITION",
    "PRECISION",
    "PREPARE",
    "PRESERVE",
    "PRIMARY",
    "PRIOR",
    "PRIVATE",
    "PRIVILEGES",
    "PROCEDURE",
    "PROCESSED",
    "PROJECT",
    "PROJECTION",
    "PROPERTY",
    "PROVISIONING",
    "PUBLIC",
    "PUT",
    "QUERY",
    "QUIT",
    "QUORUM",
    "RAISE",
    "RANDOM",
    "RANGE",
    "RANK",
    "RAW",
    "READ",
    "READS",
    "REAL",
    "REBUILD",
    "RECORD",
    "RECURSIVE",
    "REDUCE",
    "REF",
    "REFERENCE",
    "REFERENCES",
    "REFERENCING",
    "REGEXP",
    "REGION",
    "REINDEX",
    "RELATIVE",
    "RELEASE",
    "REMAINDER",
    "RENAME",
    "REPEAT",
    "REPLACE",
    "REQUEST",
    "RESET",
    "RESIGNAL",
    "RESOURCE",
    "RESPONSE",
    "RESTORE",
    "RESTRICT",
    "RESULT",
    "RETURN",
    "RETURNING",
    "RETURNS",
    "REVERSE",
    "REVOKE",
    "RIGHT",
    "ROLE",
    "ROLES",
    "ROLLBACK",
    "ROLLUP",
    "ROUTINE",
    "ROW",
    "ROWS",
    "RULE",
    "RULES",
    "SAMPLE",
    "SATISFIES",
    "SAVE",
    "SAVEPOINT",
    "SCAN",
    "SCHEMA",
    "SCOPE",
    "SCROLL",
    "SEARCH",
    "SECOND",
    "SECTION",
    "SEGMENT",
    "SEGMENTS",
    "SELECT",
    "SELF",
    "SEMI",
    "SENSITIVE",
    "SEPARATE",
    "SEQUENCE",
    "SERIALIZABLE",
    "SESSION",
    "SET",
    "SETS",
    "SHARD",
    "SHARE",
    "SHARED",
    "SHORT",
    "SHOW",
    "SIGNAL",
    "SIMILAR",
    "SIZE",
    "SKEWED",
    "SMALLINT",
    "SNAPSHOT",
    "SOME",
    "SOURCE",
    "SPACE",
    "SPACES",
    "SPARSE",
    "SPECIFIC",
    "SPECIFICTYPE",
    "SPLIT",
    "SQL",
    "SQLCODE",
    "SQLERROR",
    "SQLEXCEPTION",
    "SQLSTATE",
    "SQLWARNING",
    "START",
    "STATE",
    "STATIC",
    "STATUS",
    "STORAGE",
    "STORE",
    "STORED",
    "STREAM",
    "STRING",
    "STRUCT",
    "STYLE",
    "SUB",
    "SUBMULTISET",
    "SUBPARTITION",
    "SUBSTRING",
    "SUBTYPE",
    "SUM",
    "SUPER",
    "SYMMETRIC",
    "SYNONYM",
    "SYSTEM",
    "TABLE",
    "TABLESAMPLE",
    "TEMP",
    "TEMPORARY",
    "TERMINATED",
    "TEXT",
    "THAN",
    "THEN",
    "THROUGHPUT",
    "TIME",
    "TIMESTAMP",
    "TIMEZONE",
    "TINYINT",
    "TO",
    "TOKEN",
    "TOTAL",
    "TOUCH",
    "TRAILING",
    "TRANSACTION",
    "TRANSFORM",
    "TRANSLATE",
    "TRANSLATION",
    "TREAT",
    "TRIGGER",
    "TRIM",
    "TRUE",
    "TRUNCATE",
    "TTL",
    "TUPLE",
    "TYPE",
    "UNDER",
    "UNDO",
    "UNION",
    "UNIQUE",
    "UNIT",
    "UNKNOWN",
    "UNLOGGED",
    "UNNEST",
    "UNPROCESSED",
    "UNSIGNED",
    "UNTIL",
    "UPDATE",
    "UPPER",
    "URL",
    "USAGE",
    "USE",
    "USER",
    "USERS",
    "USING",
    "UUID",
    "VACUUM",
    "VALUE",
    "VALUED",
    "VALUES",
    "VARCHAR",
    "VARIABLE",
    "VARIANCE",
    "VARINT",
    "VARYING",
    "VIEW",
    "VIEWS",
    "VIRTUAL",
    "VOID",
    "WAIT",
    "WHEN",
    "WHENEVER",
    "WHERE",
    "WHILE",
    "WINDOW",
    "WITH",
    "WITHIN",
    "WITHOUT",
    "WORK",
    "WRAPPED",
    "WRITE",
    "YEAR",
    "ZONE",
]


class UpdateExpressionBuilder:
    def __init__(self, model_cls: Type[DatabaseModel]):
        self.SET = []
        self.ADD = []
        self.REMOVE = []
        self.DELETE = []
        self.expression_attribute_values = {}
        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 = self.safe_name(key)
            self.expression_attribute_names[safe_key] = key
            return safe_key
        return key

    def _get_field_type(self, key):
        return self.model_schema["properties"].get(key, {}).get("type", None)

    def set(self, key, value, if_not_exists=False, list_append=None):
        safe_key = self._safe_name(key)
        placeholder = ":" + key
        if if_not_exists:
            self.SET.append(f"{safe_key} = if_not_exists({safe_key}, {placeholder})")
        elif list_append is not None:
            append_placeholder = ":append_" + key
            self.SET.append(f"{safe_key} = list_append({safe_key}, {append_placeholder})")
            self.expression_attribute_values[append_placeholder] = list_append
        else:
            self.SET.append(f"{safe_key} = {placeholder}")

        self.expression_attribute_values[placeholder] = value
        return self

    def add(self, key, value):
        safe_key = self._safe_name(key)
        data_type = self._get_field_type(key)
        if "." in key:
            raise ValueError("The ADD action cannot be used on nested attributes.")
        if data_type not in ["integer", "float", "number", "array"]:
            raise ValueError(f"The ADD action does not support '{data_type}' data type.")

        placeholder = ":" + key
        self.ADD.append(f"{safe_key} {placeholder}")
        self.expression_attribute_values[placeholder] = value
        return self

    def remove(self, key):
        safe_key = self._safe_name(key)
        self.REMOVE.append(safe_key)
        self.expression_attribute_names[safe_key] = key
        return self

    def delete(self, key, value):
        data_type = self._get_field_type(key)
        safe_key = self._safe_name(key)
        if data_type != "array":
            raise ValueError(f"The DELETE action only supports set data types. Got '{data_type}' instead.")

        placeholder = ":" + key
        self.DELETE.append(f"{safe_key} {placeholder}")
        self.expression_attribute_values[placeholder] = value
        self.expression_attribute_names[safe_key] = key
        return self

    def build(self) -> Tuple[str, dict[str, Any], dict[str, Any]]:
        expression_parts = []

        if self.SET:
            expression_parts.append("SET " + ", ".join(self.SET))
        if self.ADD:
            expression_parts.append("ADD " + ", ".join(self.ADD))
        if self.REMOVE:
            expression_parts.append("REMOVE " + ", ".join(self.REMOVE))
        if self.DELETE:
            expression_parts.append("DELETE " + ", ".join(self.DELETE))

        update_expression = " ".join(expression_parts)
        return (
            update_expression,
            self.expression_attribute_values,
            self.expression_attribute_names,
        )


class DatabaseModelUpdateExpressionBuilder(UpdateExpressionBuilder):
    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.__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,
            model=self.model,
        )

src/statikk/expressions.py
terinia/statikk
from __future__ import annotations

import uuid
from typing import Optional, List, Any, Set, Type

from boto3.dynamodb.conditions import ComparisonCondition
from pydantic import BaseModel, model_serializer, model_validator
from pydantic.fields import FieldInfo
from pydantic_core._pydantic_core import PydanticUndefined

from statikk.conditions import Condition
from statikk.expressions import DatabaseModelUpdateExpressionBuilder


class Key(BaseModel):
    name: str
    type: Type = str
    default: Optional[Any] = ""


class GSI(BaseModel):
    name: str = "main-index"
    hash_key: Key
    sort_key: Key
    projection_type: str = "ALL"


class KeySchema(BaseModel):
    hash_key: str
    sort_key: Optional[str] = None


class Index(BaseModel):
    value: Optional[Any] = None
    index_names: Optional[List[str]] = ["main-index"]

    @model_serializer
    def ser_model(self) -> Any:
        return self.value


class IndexPrimaryKeyField(Index):
    pass


class IndexSecondaryKeyField(Index):
    order: Optional[int] = None


class DatabaseModel(BaseModel):
    id: str

    @classmethod
    def model_type(cls):
        return cls.__name__

    @classmethod
    def include_type_in_sort_key(cls):
        return True

    @classmethod
    def type_is_primary_key(cls):
        return False

    @classmethod
    def set_table_ref(cls, table):
        cls._table = table

    @classmethod
    def batch_write(cls):
        return cls._table.batch_write()

    @classmethod
    def query(
        cls,
        hash_key: Condition,
        range_key: Optional[Condition] = None,
        filter_condition: Optional[ComparisonCondition] = None,
        index_name: Optional[str] = None,
    ):
        return cls._table.query_index(
            hash_key=hash_key,
            model_class=cls,
            range_key=range_key,
            filter_condition=filter_condition,
            index_name=index_name,
        )

    def save(self):
        return self._table.put_item(self)

    def delete(self):
        return self._table.delete_item(self.id)

    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,
            model=model,
        )

    @classmethod
    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,
        consistent_read: bool = False,
    ):
        return cls._table.scan(model_class=cls, filter_condition=filter_condition)

    @staticmethod
    def _index_types() -> Set[Type[Index]]:
        return {Index, IndexPrimaryKeyField, IndexSecondaryKeyField}

    @staticmethod
    def _is_index_field(field: FieldInfo) -> bool:
        index_types = {Index, IndexPrimaryKeyField, IndexSecondaryKeyField}
        return field.annotation in index_types

    @classmethod
    def _create_index_field_from_shorthand(cls, field: FieldInfo, value: str) -> FieldInfo:
        annotation = field.annotation
        extra_fields = dict()
        if field.default is not PydanticUndefined:
            extra_fields["index_names"] = field.default.index_names
            if field.annotation is IndexSecondaryKeyField:
                extra_fields["order"] = field.default.order
        return annotation(value=value, **extra_fields)

    @model_validator(mode="before")
    @classmethod
    def check_boxed_indexes(cls, data: Any) -> Any:
        """
        This method allows for a cleaner interface when working with model indexes.
        Instead of having to manually box the index value, this method will do it for you.
        For example:
          my_model = MyModel(id="123", my_index="abc")
        Is equivalent to:
          my_model = MyModel(id="123", my_index=IndexPrimaryKeyField(value="abc"))
        """
        if isinstance(data, dict):
            if "id" not in data:
                data["id"] = str(uuid.uuid4())

            for key, value in data.items():
                field = cls.model_fields[key]
                if cls._is_index_field(field) and not isinstance(value, dict):
                    if isinstance(value, tuple(cls._index_types())):
                        continue
                    else:
                        data[key] = cls._create_index_field_from_shorthand(field, value)
        return data

    @model_validator(mode="after")
    def validate_index_fields(self):
        sort_key_field_orders = [
            getattr(self, field_name).order
            for field_name, field_info in self.model_fields.items()
            if field_info.annotation is not None
            if field_info.annotation is IndexSecondaryKeyField
        ]
        have_orders_defined = any([order for order in sort_key_field_orders])
        all_orders_defined = all([order for order in sort_key_field_orders])
        use_default_ordering = all([order is None for order in sort_key_field_orders])
        if have_orders_defined and not all_orders_defined and not use_default_ordering:
            raise ValueError(
                f"`order` is not defined on at least one of the Index keys on model {type(self)}. Please set the order for all sort key fields."
            )
        if all_orders_defined and len(set(sort_key_field_orders)) != len(sort_key_field_orders):
            raise ValueError(
                f"Duplicate `order` values found on model {type(self)}. Please ensure that all `order` values are unique."
            )

src/statikk/models.py
terinia/statikk
"""
    Dummy conftest.py for statikk.

    If you don't know what this is for, just leave it empty.
    Read more about conftest.py under:
    - https://docs.pytest.org/en/stable/fixture.html
    - https://docs.pytest.org/en/stable/writing_plugins.html
"""

# import pytest

tests/conftest.py
terinia/statikk
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

from statikk.conditions import Equals, BeginsWith
from statikk.engine import (
    Table,
    InvalidIndexNameError,
    IncorrectSortKeyError,
    ItemNotFoundError,
)
from statikk.models import (
    DatabaseModel,
    IndexPrimaryKeyField,
    IndexSecondaryKeyField,
    KeySchema,
    GSI,
    Key,
)


class MyAwesomeModel(DatabaseModel):
    player_id: IndexPrimaryKeyField
    tier: IndexSecondaryKeyField
    name: str = "Foo"
    values: set = {1, 2, 3, 4}
    cost: int = 4
    probability: float = 0.5


class SimpleModel(DatabaseModel):
    player_id: IndexPrimaryKeyField
    board_id: IndexSecondaryKeyField


class DoubleIndexModel(DatabaseModel):
    player_id: IndexPrimaryKeyField
    tier: IndexSecondaryKeyField
    card_template_id: IndexPrimaryKeyField = IndexPrimaryKeyField(index_names=["secondary-index"])
    added_at: IndexSecondaryKeyField = IndexSecondaryKeyField(index_names=["secondary-index"])


class MultiIndexModel(DatabaseModel):
    player_id: IndexPrimaryKeyField

    card_template_id: IndexPrimaryKeyField = IndexPrimaryKeyField(index_names=["secondary-index"])
    tier: IndexSecondaryKeyField = IndexSecondaryKeyField(index_names=["secondary-index", "main-index"])
    values: List[int] = [1, 2, 3, 4]

    def include_type_in_sort_key(cls):
        return False


class SomeOtherIndexModel(DatabaseModel):
    player_id: IndexPrimaryKeyField = IndexPrimaryKeyField(index_names=["my-awesome-index"])
    tier: IndexSecondaryKeyField = IndexSecondaryKeyField(index_names=["my-awesome-index"])


def _create_dynamodb_table(table):
    table.create(aws_region="eu-west-1")


def test_create_my_awesome_model():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY")
    table.put_item(model)
    assert table.get_item(id=model.id, model_class=MyAwesomeModel).model_dump() == {
        "id": "foo",
        "player_id": "123",
        "tier": "LEGENDARY",
        "gsi_pk": "123",
        "gsi_sk": "MyAwesomeModel|LEGENDARY",
        "name": "Foo",
        "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)
    assert table.get_item(id=model_2.id, model_class=MyAwesomeModel).model_dump() == {
        "id": "foo-2",
        "player_id": "123",
        "tier": "EPIC",
        "gsi_pk": "123",
        "gsi_sk": "MyAwesomeModel|EPIC",
        "name": "FooFoo",
        "values": {1, 2, 3, 4},
        "cost": 4,
        "type": "MyAwesomeModel",
        "probability": 0.5,
    }
    mock_dynamodb().stop()


def test_multi_index_table():
    mock_dynamodb().start()
    table = Table(
        name="my-table",
        key_schema=KeySchema(hash_key="id"),
        indexes=[
            GSI(
                name="main-index",
                hash_key=Key(name="gsi_pk"),
                sort_key=Key(name="gsi_sk"),
            ),
            GSI(
                name="secondary-index",
                hash_key=Key(name="gsi_pk_2"),
                sort_key=Key(name="gsi_sk_2", type=datetime),
            ),
        ],
        models=[DoubleIndexModel],
    )
    _create_dynamodb_table(table)
    my_model = DoubleIndexModel(
        id="foo",
        player_id="123",
        tier="LEGENDARY",
        card_template_id="abc",
        added_at=datetime(2023, 9, 10, 12, 0, 0),
    )
    table.put_item(my_model)
    assert table.get_item(id=my_model.id, model_class=DoubleIndexModel).model_dump() == {
        "id": "foo",
        "player_id": "123",
        "tier": "LEGENDARY",
        "card_template_id": "abc",
        "added_at": 1694340000,
        "gsi_pk": "123",
        "gsi_sk": "DoubleIndexModel|LEGENDARY",
        "gsi_pk_2": "abc",
        "gsi_sk_2": datetime(2023, 9, 10, 12, 0),
        "type": "DoubleIndexModel",
    }
    mock_dynamodb().stop()


def test_incorrect_index_type():
    mock_dynamodb().start()
    table = Table(
        name="my-table",
        key_schema=KeySchema(hash_key="id"),
        indexes=[
            GSI(
                name="main-index",
                hash_key=Key(name="gsi_pk"),
                sort_key=Key(name="gsi_sk"),
            ),
            GSI(
                name="secondary-index",
                hash_key=Key(name="gsi_pk_2"),
                sort_key=Key(name="gsi_sk_2", type=datetime),
            ),
        ],
        models=[DoubleIndexModel],
    )
    _create_dynamodb_table(table)
    my_model = DoubleIndexModel(
        id="foo",
        player_id="123",
        tier="LEGENDARY",
        card_template_id="abc",
        added_at="2023-01-01 12:00:00",
    )

    with pytest.raises(IncorrectSortKeyError) as e:
        table.put_item(my_model)
    assert (
        e.value.args[0]
        == f"Incorrect sort key type. Sort key type for sort key 'gsi_sk_2' should be: <class 'datetime.datetime'> but got: <class 'str'>"
    )


def test_multi_field_index():
    mock_dynamodb().start()
    table = Table(
        name="my-table",
        key_schema=KeySchema(hash_key="id"),
        indexes=[
            GSI(
                name="main-index",
                hash_key=Key(name="gsi_pk"),
                sort_key=Key(name="gsi_sk"),
            ),
            GSI(
                name="secondary-index",
                hash_key=Key(name="gsi_pk_2"),
                sort_key=Key(name="gsi_sk_2"),
            ),
        ],
        models=[MultiIndexModel],
    )
    _create_dynamodb_table(table)
    model = MultiIndexModel(id="card-id", player_id="123", card_template_id="abc", tier="LEGENDARY")
    table.put_item(model)
    assert table.get_item(id=model.id, model_class=MultiIndexModel).model_dump() == {
        "card_template_id": "abc",
        "gsi_pk": "123",
        "gsi_pk_2": "abc",
        "gsi_sk": "LEGENDARY",
        "gsi_sk_2": "LEGENDARY",
        "id": "card-id",
        "player_id": "123",
        "tier": "LEGENDARY",
        "values": [1, 2, 3, 4],
        "type": "MultiIndexModel",
    }
    mock_dynamodb().stop()


def test_integration_get_item():
    mock_dynamodb().start()
    table = Table(
        name="my-table",
        key_schema=KeySchema(hash_key="id"),
        indexes=[
            GSI(
                name="main-index",
                hash_key=Key(name="gsi_pk"),
                sort_key=Key(name="gsi_sk"),
            ),
            GSI(
                name="secondary-index",
                hash_key=Key(name="gsi_pk_2"),
                sort_key=Key(name="gsi_sk_2"),
            ),
        ],
        models=[MultiIndexModel],
    )
    _create_dynamodb_table(table)
    model = MultiIndexModel(id="card-id", player_id="123", card_template_id="abc", tier="LEGENDARY")
    table.put_item(model)
    item = table.get_item("card-id", MultiIndexModel)
    assert item.id == model.id
    assert item.player_id == model.player_id
    assert item.card_template_id == model.card_template_id
    assert item.tier == model.tier
    assert item.gsi_pk == "123"
    assert item.gsi_pk_2 == "abc"
    assert item.gsi_sk == "LEGENDARY"
    assert item.gsi_sk_2 == "LEGENDARY"
    mock_dynamodb().stop()


def test_query_model_index():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY", name="Terror From Below")
    table.put_item(model)
    model_2 = MyAwesomeModel(id="foo-2", player_id="123", tier="EPIC")
    table.put_item(model_2)
    models = list(
        table.query_index(
            hash_key=Equals("123"),
            range_key=BeginsWith("LEG"),
            filter_condition=Attr("name").eq("Terror From Below"),
            model_class=MyAwesomeModel,
        )
    )
    assert len(models) == 1
    assert models[0].id == model.id
    assert models[0].model_type == model.model_type
    assert models[0].tier == model.tier
    mock_dynamodb().stop()


def test_query_index_name_is_provided():
    mock_dynamodb().start()
    table = Table(
        name="my-dynamodb-table",
        key_schema=KeySchema(hash_key="id"),
        indexes=[
            GSI(
                name="my-awesome-index",
                hash_key=Key(name="gsi_pk"),
                sort_key=Key(name="gsi_sk"),
            )
        ],
        models=[SomeOtherIndexModel],
    )
    _create_dynamodb_table(table)
    model = SomeOtherIndexModel(id="foo", player_id="123", tier="LEGENDARY")
    table.put_item(model)
    model_2 = SomeOtherIndexModel(id="foo-2", player_id="123", tier="EPIC")
    table.put_item(model_2)
    models = list(
        table.query_index(
            index_name="my-awesome-index",
            hash_key=Equals("123"),
            range_key=BeginsWith("SomeOtherIndexModel"),
            filter_condition=Attr("tier").eq("LEGENDARY"),
            model_class=SomeOtherIndexModel,
        )
    )
    assert len(models) == 1
    assert models[0].id == model.id
    assert models[0].model_type == model.model_type
    assert models[0].tier == model.tier
    mock_dynamodb().stop()


def test_batch_get_items():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY")
    model_2 = MyAwesomeModel(id="foo-2", player_id="123", tier="LEGENDARY")
    table.put_item(model)
    table.put_item(model_2)
    models = table.batch_get_items(["foo", "foo-2"], MyAwesomeModel, batch_size=1)
    assert len(models) == 2
    assert models[0].id == model.id
    assert models[0].model_type == model.model_type
    assert models[0].tier == model.tier
    assert models[1].id == model_2.id
    assert models[1].model_type == model_2.model_type
    assert models[1].tier == model_2.tier
    mock_dynamodb().stop()


def test_batch_write():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)

    with table.batch_write() as batch:
        for i in range(30):
            model = MyAwesomeModel(id=f"foo_{i}", player_id="123", tier="LEGENDARY")
            batch.put(model)

    models = list(
        table.query_index(
            index_name="main-index",
            hash_key=Equals("123"),
            range_key=BeginsWith("MyAwesomeModel"),
            filter_condition=Attr("tier").eq("LEGENDARY"),
            model_class=MyAwesomeModel,
        )
    )
    assert len(models) == 30

    with table.batch_write() as batch:
        for model in models:
            batch.delete(model)

    models = list(
        table.query_index(
            index_name="main-index",
            hash_key=Equals("123"),
            range_key=BeginsWith("MyAwesomeModel"),
            filter_condition=Attr("tier").eq("LEGENDARY"),
            model_class=MyAwesomeModel,
        )
    )
    assert len(models) == 0

    mock_dynamodb().stop()


def test_query_index_does_not_exist():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    with pytest.raises(InvalidIndexNameError) as e:
        list(
            table.query_index(
                hash_key=Equals("123"),
                range_key=BeginsWith("foo"),
                index_name="does-not-exist",
                model_class=MyAwesomeModel,
            )
        )
    assert e.value.args[0] == "The provided index name 'does-not-exist' is not configured on the table."
    mock_dynamodb().stop()


def test_table_delegates():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY").save()
    saved_model = MyAwesomeModel.get(model.id)
    assert model == saved_model
    models = list(MyAwesomeModel.query(hash_key=Equals("123")))
    assert len(models) == 1
    assert models[0] == model
    model_2 = MyAwesomeModel(id="foo-2", player_id="123", tier="bar")
    model_3 = MyAwesomeModel(id="foo-3", player_id="123", tier="bar")
    with MyAwesomeModel.batch_write() as batch:
        batch.put(model_2)
        batch.put(model_3)

    saved_models = MyAwesomeModel.batch_get(
        ["foo", "foo-2", "foo-3"],
    )
    assert len(saved_models) == 3
    assert saved_models[0] == model
    assert saved_models[1] == model_2
    assert saved_models[2] == model_3


def test_exclude_type_from_sort_key():
    class ExcludeTypeModel(DatabaseModel):
        player_id: IndexPrimaryKeyField
        tier: IndexSecondaryKeyField

        @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=[ExcludeTypeModel],
    )
    _create_dynamodb_table(table)
    model = ExcludeTypeModel(id="foo", player_id="123", tier="LEGENDARY").save()
    saved_model = ExcludeTypeModel.get(model.id)
    assert "ExcludeTypeModel" not in saved_model.gsi_sk
    mock_dynamodb().stop()


def test_type_is_primary_key():
    class TypeIsPrimaryKeyModel(DatabaseModel):
        tier: IndexSecondaryKeyField = IndexSecondaryKeyField(index_names=["main-index", "secondary-index"])
        foo: IndexPrimaryKeyField = IndexPrimaryKeyField(index_names=["secondary-index"])

        @classmethod
        def type_is_primary_key(cls):
            return True

        @classmethod
        def include_type_in_sort_key(cls):
            return False

        @classmethod
        def model_type(cls):
            return "my-type"

    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"),
            ),
            GSI(
                name="secondary-index",
                hash_key=Key(name="gsi_pk_2"),
                sort_key=Key(name="gsi_sk_2"),
            ),
        ],
        models=[TypeIsPrimaryKeyModel],
    )
    _create_dynamodb_table(table)
    model = TypeIsPrimaryKeyModel(tier="LEGENDARY", foo="Bar").save()
    assert model.gsi_pk == "my-type"
    assert model.gsi_sk == "LEGENDARY"
    assert model.gsi_pk_2 == "Bar"
    assert model.gsi_sk_2 == "LEGENDARY"
    mock_dynamodb().stop()


def test_delete_model():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY")
    model_2 = MyAwesomeModel(id="foo-2", player_id="123", tier="EPIC", name="FooFoo")
    model_3 = MyAwesomeModel(id="foo-3", player_id="123", tier="EPIC", name="FooFooFoo")
    table.put_item(model)
    table.put_item(model_2)
    table.delete_item(model.id)
    model_3.delete()
    assert list(table.query_index("123", MyAwesomeModel)) == [model_2]

    mock_dynamodb().stop()


def test_get_item_does_not_exist():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)

    with pytest.raises(ItemNotFoundError) as e:
        table.get_item("foo", MyAwesomeModel)

    mock_dynamodb().stop()


def test_update():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY", name="FooFoo", values={1, 2, 3, 4})
    model.save()
    (
        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
    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()


def test_scan():
    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=[MyAwesomeModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY", name="FooFoo", values={1, 2, 3, 4})
    model_2 = MyAwesomeModel(id="foo-2", player_id="123", tier="EPIC", name="BarBar")
    model.save()
    model_2.save()
    items = list(MyAwesomeModel.scan())
    assert len(items) == 2
    mock_dynamodb().stop()


def test_query_no_range_key_provided():
    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=[MyAwesomeModel, SimpleModel],
    )
    _create_dynamodb_table(table)
    model = MyAwesomeModel(id="foo", player_id="123", tier="LEGENDARY", name="FooFoo", values={1, 2, 3, 4})
    model.save()
    model_2 = SimpleModel(
        player_id="123",
        board_id="456",
    )
    model_2.save()

    my_awesome_models = list(MyAwesomeModel.query(hash_key=Equals("123")))
    assert len(my_awesome_models) == 1
    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
        unit_class: IndexSecondaryKeyField = IndexSecondaryKeyField(order=2)
        tier: IndexSecondaryKeyField = IndexSecondaryKeyField(order=1)
        name: str = "Foo"
        values: set = {1, 2, 3, 4}
        cost: int = 4

    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=[ModelWithIndexOrdersDefined],
    )
    _create_dynamodb_table(table)
    model = ModelWithIndexOrdersDefined(id="123", player_id="456", unit_class="Mage", tier="EPIC")
    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",
    }

tests/test_engine.py
terinia/statikk
import boto3
import pytest
from moto import mock_dynamodb

from statikk.expressions import (
    UpdateExpressionBuilder,
)
from statikk.models import DatabaseModel


class MyModel(DatabaseModel):
    attribute1: int
    attribute2: list
    attribute3: str
    attribute4: set
    name: str  # reserved keyword
    value: int  # reserved keyword
    values: set  # reserved keyword


# A fixture to create a mock DynamoDB table
@pytest.fixture
def create_mocked_table():
    with mock_dynamodb():
        dynamodb = boto3.resource("dynamodb", region_name="us-west-2")
        dynamodb.create_table(
            TableName="TestTable",
            KeySchema=[{"AttributeName": "id", "KeyType": "HASH"}],
            AttributeDefinitions=[{"AttributeName": "id", "AttributeType": "S"}],
            ProvisionedThroughput={"ReadCapacityUnits": 1, "WriteCapacityUnits": 1},
        )
        yield


def test_set_method(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    builder.set("attribute1", 10).set("name", "Bob")
    expr, values, attribute_names = builder.build()
    assert expr == "SET attribute1 = :attribute1, #n_name = :name"
    assert values == {":attribute1": 10, ":name": "Bob"}
    assert attribute_names == {"#n_name": "name"}


def test_add_method_valid(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    builder.add("value", 5).add("attribute1", 2)
    expr, values, attr_names = builder.build()
    assert expr == "ADD #n_value :value, attribute1 :attribute1"
    assert values == {":value": 5, ":attribute1": 2}
    assert attr_names == {"#n_value": "value"}


def test_add_method_invalid_type(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    with pytest.raises(ValueError, match="The ADD action does not support"):
        builder.add("attribute3", "invalid")


def test_add_method_invalid_nested(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    with pytest.raises(ValueError, match="The ADD action cannot be used on nested attributes."):
        builder.add("attribute.nested", 5)


def test_remove_method(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    builder.remove("attribute3").remove("name")
    expr, _, attr_names = builder.build()
    assert expr == "REMOVE attribute3, #n_name"
    assert attr_names == {"#n_name": "name", "attribute3": "attribute3"}


def test_delete_method_valid(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    builder.delete("attribute4", {"value"}).delete("values", {1})
    expr, values, attr_names = builder.build()
    assert expr == "DELETE attribute4 :attribute4, #n_values :values"
    assert values == {":attribute4": {"value"}, ":values": {1}}
    assert attr_names == {"#n_values": "values", "attribute4": "attribute4"}


def test_delete_method_invalid(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    with pytest.raises(ValueError, match="The DELETE action only supports"):
        builder.delete("attribute1", 5)


def test_multiple_methods(create_mocked_table):
    builder = UpdateExpressionBuilder(MyModel)
    builder.set("attribute1", 10).add("attribute1", 5).remove("attribute3").delete("attribute4", {"value"})
    expr, values, attr_names = builder.build()
    assert (
        expr
        == "SET attribute1 = :attribute1 ADD attribute1 :attribute1 REMOVE attribute3 DELETE attribute4 :attribute4"
    )
    assert values == {":attribute1": 5, ":attribute4": {"value"}}

tests/test_expressions.py
terinia/statikk
from unittest.mock import patch

import pytest
from pydantic_core._pydantic_core import ValidationError

from statikk.models import (
    DatabaseModel,
    IndexPrimaryKeyField,
    IndexSecondaryKeyField,
)


def test_index_shorthand():
    class Foo(DatabaseModel):
        player_id: IndexSecondaryKeyField

        @classmethod
        def type_is_primary_key(cls):
            return True

    with patch("statikk.models.uuid.uuid4", return_value="123"):
        foo = Foo(player_id="abc")
        assert foo.model_type() == "Foo"
        assert foo.id == "123"
        assert foo.player_id.value == "abc"
        assert foo.player_id.index_names == ["main-index"]


def test_index_configuration():
    class FooNotMainIndex(DatabaseModel):
        type: IndexPrimaryKeyField = IndexPrimaryKeyField(index_names=["not-main-index"])

    foo = FooNotMainIndex(type="FooNotMainIndex")
    assert foo.type.value == "FooNotMainIndex"
    assert foo.type.index_names == ["not-main-index"]


def test_invalid_order_config_on_indexes():
    class InvalidOrderingModel(DatabaseModel):
        player_id: IndexPrimaryKeyField
        tier: IndexSecondaryKeyField = IndexSecondaryKeyField(order=1)
        unit_class: IndexSecondaryKeyField

    with pytest.raises(ValidationError):
        InvalidOrderingModel(player_id="abc", unit_class="foo")


def test_all_orders_are_configued_on_indexes():
    class AllOrdersConfiguredModel(DatabaseModel):
        player_id: IndexPrimaryKeyField
        tier: IndexSecondaryKeyField = IndexSecondaryKeyField(order=1)
        unit_class: IndexSecondaryKeyField = IndexSecondaryKeyField(order=2)

    model = AllOrdersConfiguredModel(player_id="abc", unit_class="foo", tier="bar")
    assert model.tier.order == 1
    assert model.unit_class.order == 2


def test_same_order_defined_on_fields():
    class SameOrderDefinedOnFields(DatabaseModel):
        player_id: IndexPrimaryKeyField
        tier: IndexSecondaryKeyField = IndexSecondaryKeyField(order=1)
        unit_class: IndexSecondaryKeyField = IndexSecondaryKeyField(order=1)

    with pytest.raises(ValidationError):
        SameOrderDefinedOnFields(player_id="abc", unit_class="foo", tier="bar")

tests/test_models.py
terinia/statikk