import json
import logging
from collections.abc import Callable
from copy import deepcopy
from datetime import UTC, datetime
from enum import StrEnum, auto
from typing import Any
from da_vinci.core.orm.orm_exceptions import MissingTableObjectAttributeError
[docs]
class TableObjectAttributeType(StrEnum):
STRING = auto()
NUMBER = auto()
BOOLEAN = auto()
DATETIME = auto()
JSON = auto() # Not safe for storing empty attributes, native
JSON_STRING = auto() # Safe for storing empty attributes
STRING_LIST = auto()
NUMBER_LIST = auto()
JSON_LIST = auto() # Not safe for storing empty attributes, native
JSON_STRING_LIST = auto() # Safe for storing empty attributes
COMPOSITE_STRING = auto() # String in DynamoDB but a tuple of strings in Python
STRING_SET = auto()
NUMBER_SET = auto()
[docs]
@classmethod
def is_list(cls, attribute_type: "TableObjectAttributeType") -> bool:
"""
Check if the attribute type is a list
Keyword Arguments:
attribute_type -- Attribute type to check
Returns:
bool
"""
return attribute_type in (cls.STRING_LIST, cls.NUMBER_LIST, cls.JSON_LIST)
[docs]
def to_str(self) -> str:
"""
Convert the attribute type to a string
Returns:
str
"""
return self.name
[docs]
class TableObjectAttribute:
[docs]
def __init__(
self,
name: str,
attribute_type: TableObjectAttributeType,
argument_names: list[str] | None = None,
custom_exporter: Callable | None = None,
custom_importer: Callable | None = None,
description: str | None = None,
dynamodb_key_name: str | None = None,
default: Any | None = None,
exclude_from_dict: bool | None = False,
exclude_from_schema_description: bool | None = False,
is_indexed: bool = True,
optional: bool = False,
) -> None:
"""
Object representing an attribute of a TableObject
Keyword Arguments:
name -- Name of the attribute
attribute_type -- Type of the attribute
argument_names -- The names of Python arguments that are merged into a composite
string. This is required when the attribute type is COMPOSITE_STRING.
custom_exporter -- Custom exporter function, called whenever data is converted for DynamoDB
custom_importer -- Custom importer function, called whenever data is loaded from DynamoDB
description -- Description of the attribute, annotation primarily
used for LLM context.
dynamodb_key_name -- Name of the DynamoDB key, defaults to the ORM naming convention based
on the attribute name
default -- Default value for the attribute, attribute is considered optional when this is set.
Accepts a value or a callable that returns a value.
exclude_from_dict -- Attribute is not added to a resulting Dict when calling to_dict()
exclude_from_schema_description -- Attribute is not included in the table object schema description
is_indexed -- Whether the attribute is able to be used to query with, defaults to True
optional -- Whether the attribute optional, defaults to False unless a default is provided
"""
self.name = name
self.description = description
self.attribute_type = attribute_type
self.exclude_from_dict = exclude_from_dict
self.exclude_from_schema_description = exclude_from_schema_description
self.is_indexed = is_indexed
if (
self.attribute_type is TableObjectAttributeType.JSON_STRING
or self.attribute_type is TableObjectAttributeType.JSON_STRING_LIST
):
self.is_indexed = False
if dynamodb_key_name:
self.dynamodb_key_name = dynamodb_key_name
else:
self.dynamodb_key_name = self.default_dynamodb_key_name(self.name)
self.argument_names = argument_names
if (
self.attribute_type == TableObjectAttributeType.COMPOSITE_STRING
and not self.argument_names
):
raise ValueError(
"argument_names must be provided when attribute_type is COMPOSITE_STRING"
)
self._default = default
if self._default is None:
self.optional = optional
else:
self.optional = True
self.custom_exporter = custom_exporter
self.custom_importer = custom_importer
[docs]
@staticmethod
def composite_string_value(values: list[str]) -> str:
"""
Return a full composite string value given a list of attribute values
Keyword Arguments:
values -- List of values to join into the full composite string value
"""
return "-".join(values)
[docs]
@staticmethod
def default_dynamodb_key_name(name: str) -> str:
"""
Convert a name to a DynamoDB key name
Keyword Arguments:
name -- Name to convert
Returns:
str
"""
return "".join([wrd.capitalize() for wrd in name.split("_")])
[docs]
@staticmethod
def timestamp_to_datetime(timestamp: int | float) -> datetime:
"""
Convert a timestamp string to a datetime
Keyword Arguments:
timestamp -- Timestamp string
Returns:
datetime
"""
return datetime.fromtimestamp(timestamp)
[docs]
@staticmethod
def datetime_to_timestamp(dt: datetime) -> float:
"""
Convert a datetime to a timestamp
Keyword Arguments:
dt -- Datetime
Returns:
float
"""
return dt.timestamp()
[docs]
def schema_to_str(self) -> str:
"""
Describe the schema for the attribute
Returns:
str
"""
descr = f"{self.name} - type: {self.attribute_type.to_str()}"
if self.description:
descr += f" description: {self.description}"
return descr
@property
def default(self) -> Any:
if callable(self._default):
return self._default()
return self._default
@property
def dynamodb_type_label(self) -> str:
"""
Get the DynamoDB type label for the attribute type
Returns:
str
"""
dynamodb_type_label = "S"
# Handle number and datetime types
if (
self.attribute_type is TableObjectAttributeType.NUMBER
or self.attribute_type is TableObjectAttributeType.DATETIME
):
dynamodb_type_label = "N"
# Handle JSON types
elif self.attribute_type is TableObjectAttributeType.JSON:
dynamodb_type_label = "M"
# Handle boolean types
elif self.attribute_type is TableObjectAttributeType.BOOLEAN:
dynamodb_type_label = "BOOL"
# Handle list types
elif TableObjectAttributeType.is_list(self.attribute_type):
dynamodb_type_label = "L"
# Handle set types
elif self.attribute_type in (
TableObjectAttributeType.STRING_SET,
TableObjectAttributeType.NUMBER_SET,
):
dynamodb_type_label = (
"SS" if self.attribute_type == TableObjectAttributeType.STRING_SET else "NS"
)
return dynamodb_type_label
def _infer_dynamodb_value(self, value: Any) -> dict:
"""
Helper method to infer DynamoDB value type for nested structures.
Keyword Arguments:
value -- Value to infer
"""
if isinstance(value, str):
return {"S": value}
elif isinstance(value, bool):
return {"BOOL": value}
elif isinstance(value, (int, float)):
return {"N": str(value)}
elif isinstance(value, dict):
if "M" in value:
return value
return {"M": {k: self._infer_dynamodb_value(v) for k, v in value.items()}}
elif isinstance(value, list):
return {"L": [self._infer_dynamodb_value(v) for v in value]}
elif value is None:
return {"NULL": True}
else:
raise ValueError(f"Unsupported value type: {type(value)}")
[docs]
def dynamodb_value(self, value: Any) -> Any:
"""
Convert a value to a DynamoDB supported value
Keyword Arguments:
value -- Value to convert
Returns:
Any
"""
if self.custom_exporter:
return self.custom_exporter(value)
# Handle number types
if self.attribute_type is TableObjectAttributeType.NUMBER:
return str(value)
# Handle datetime types
elif self.attribute_type is TableObjectAttributeType.DATETIME:
if not value:
return str(0)
return str(float(self.datetime_to_timestamp(value)))
# Handle JSON types
elif self.attribute_type is TableObjectAttributeType.JSON:
if isinstance(value, str):
value = json.loads(value)
elif not value:
return None
return {k: self._infer_dynamodb_value(v) for k, v in value.items()}
elif (
self.attribute_type is TableObjectAttributeType.JSON_STRING
or self.attribute_type is TableObjectAttributeType.JSON_STRING_LIST
):
if not value:
if self.attribute_type is TableObjectAttributeType.JSON_STRING_LIST:
return "[]"
else:
return "{}"
return json.dumps(value)
# Handle composite string types
elif self.attribute_type is TableObjectAttributeType.COMPOSITE_STRING:
if isinstance(value, str):
return value
return TableObjectAttribute.composite_string_value(value)
# Handle list types
elif TableObjectAttributeType.is_list(self.attribute_type):
# Specifically handle JSON_LIST
if self.attribute_type is TableObjectAttributeType.JSON_LIST:
if not value:
return None
# Ensure each element in the list is converted properly
return [
{"M": json.loads(item) if isinstance(item, str) else item} for item in value
]
if not value:
return []
label = "N" if self.attribute_type is TableObjectAttributeType.NUMBER_LIST else "S"
return [{label: str(val)} for val in value]
# Handle string set types
elif self.attribute_type == TableObjectAttributeType.STRING_SET:
if not value:
return None
return list(value) # DynamoDB stores sets as lists in JSON format
# Handle number set types
elif self.attribute_type == TableObjectAttributeType.NUMBER_SET:
if not value:
return None
return [str(val) for val in value]
# Handle boolean types
elif not isinstance(value, bool) and not value:
return str(value)
return value
[docs]
def as_dynamodb_attribute(self, value: Any) -> dict | None:
"""
Return the attribute as a DynamoDB attribute
Keyword Arguments:
value -- Value to convert
"""
# Skip None values or empty sets/dictionaries for JSON and Set types
if self.attribute_type in (
TableObjectAttributeType.STRING_SET,
TableObjectAttributeType.NUMBER_SET,
) and (value is None or not value):
return None # Skip empty sets
if self.attribute_type in (
TableObjectAttributeType.JSON,
TableObjectAttributeType.JSON_LIST,
) and (value is None or (isinstance(value, dict) and not value)):
return None # Skip empty JSON or JSON_LIST
return {
self.dynamodb_key_name: {
self.dynamodb_type_label: self.dynamodb_value(value),
}
}
def _infer_python_value(self, value: dict) -> Any:
"""
Helper method to convert DynamoDB types back to Python values.
Keyword Arguments:
value -- Value to convert
"""
if "S" in value:
return value["S"]
elif "N" in value:
return float(value["N"]) if "." in value["N"] else int(value["N"])
elif "BOOL" in value:
return value["BOOL"]
elif "M" in value:
return {k: self._infer_python_value(v) for k, v in value["M"].items()}
elif "L" in value:
return [self._infer_python_value(v) for v in value["L"]]
elif "NULL" in value:
return None
elif "SS" in value:
return set(value["SS"])
elif "NS" in value:
return set(map(int, value["NS"]))
else:
raise ValueError(f"Unsupported DynamoDB value type: {value}")
[docs]
def true_value(self, value: Any) -> Any:
"""
Return the attribute value as a Python value
"""
value = value[self.dynamodb_type_label]
if isinstance(value, str) and value == "None":
return None
if self.attribute_type is TableObjectAttributeType.NUMBER:
if "." in value:
return float(value)
return int(value)
elif self.attribute_type is TableObjectAttributeType.DATETIME:
if float(value) == 0.0:
return None
return self.timestamp_to_datetime(float(value))
# Handle JSON_LIST
elif self.attribute_type is TableObjectAttributeType.JSON_LIST:
# Convert each item in the list from DynamoDB format to a Python dictionary
return [self._infer_python_value(item) for item in value]
# Handle other list types
elif TableObjectAttributeType.is_list(self.attribute_type):
label = "N" if self.attribute_type is TableObjectAttributeType.NUMBER_LIST else "S"
return [item[label] for item in value]
elif self.attribute_type == TableObjectAttributeType.STRING_SET:
return set(value) # Convert list back to set
elif self.attribute_type == TableObjectAttributeType.NUMBER_SET:
return set(value)
elif self.attribute_type is TableObjectAttributeType.COMPOSITE_STRING:
return tuple(value.split("-"))
elif self.attribute_type is TableObjectAttributeType.JSON:
# If the value is already a dict (DynamoDB MAP), return it as is
if isinstance(value, dict):
return {k: self._infer_python_value(v) for k, v in value.items()}
elif (
self.attribute_type is TableObjectAttributeType.JSON_STRING
or self.attribute_type is TableObjectAttributeType.JSON_STRING_LIST
):
if not value:
if self.attribute_type is TableObjectAttributeType.JSON_STRING_LIST:
return []
else:
return {}
return json.loads(value)
return value
[docs]
def set_attribute(self, obj: Any, value: Any) -> None:
"""
Set the attribute on an object
Keyword Arguments:
obj -- Object to set the attribute on
value -- Value to set
"""
setattr(obj, self.name, value)
[docs]
def from_dynamodb_attribute(self, value: Any) -> Any:
"""
Return the attribute value as a Python value, run a
custom importer if one was set
"""
true_val = self.true_value(value)
if self.custom_importer:
return self.custom_importer(true_val)
return true_val
[docs]
class TableObject:
"""
Base class for Table object definitions
Class Attributes:
attribute_lookup_prefix: Attribute lookup prefix, prefixes the attribute name when retrieving attributes
attributes: List of attributes
description: Description of the table
object_name: Name of the object, defaults to the class name. This should be set when
dynamically defining table objects.
partition_key_attribute: Partition key attribute
sort_key_attribute: Sort key attribute
table_name: Name of the table
ttl_attribute: Optional TTL attribute
Example:
```
from uuid import uuid4
from da_vinci.core.orm.table_object import (
TableObject,
TableObjectAttribute,
TableObjectAttributeType,
)
class MyTableObject(TableObject):
partition_key_attribute = TableObjectAttribute(
name='my_pk',
attribute_type=TableObjectAttributeType.STRING,
)
sort_key_attribute = TableObjectAttribute(
name='my_sk',
attribute_type=TableObjectAttributeType.STRING,
default=lambda: str(uuid4()),
)
table_name = 'my_table'
attributes = [
TableObjectAttribute(
name='created_on',
attribute_type=TableObjectAttributeType.DATETIME,
default=lambda: datetime.now(),
),
]
```
"""
partition_key_attribute: TableObjectAttribute
table_name: str
attribute_lookup_prefix: str | None = None
attributes: list[TableObjectAttribute] = []
description: str | None = None
object_name: str | None = None
sort_key_attribute: TableObjectAttribute | None = None
ttl_attribute: TableObjectAttribute | None = None
[docs]
def __init__(self, **kwargs) -> None:
"""
Base class for Table objects
"""
self.__attr_index__ = {}
for attr in self.all_attributes():
self.__attr_index__[attr.name] = attr
if (
attr.attribute_type is TableObjectAttributeType.COMPOSITE_STRING
and attr.argument_names
and set(kwargs.keys()).issuperset(set(attr.argument_names))
):
composite_args: list = []
for arg in attr.argument_names:
composite_args.append(kwargs[arg])
attr.set_attribute(self, attr.composite_string_value(composite_args))
elif attr.name in kwargs:
# If the value is None and the attribute has a default, use the default
if not kwargs[attr.name] and attr.default:
attr.set_attribute(self, attr.default)
else:
attr.set_attribute(self, kwargs[attr.name])
else:
# If the attribute is optional set default, default will either be a value,
# a callable, or None which is fine for optional attributes
if attr.optional:
attr.set_attribute(self, attr.default)
else:
raise MissingTableObjectAttributeError(attr.name)
[docs]
@classmethod
def define(
cls,
partition_key_attribute: TableObjectAttribute,
object_name: str,
table_name: str,
attribute_lookup_prefix: str | None = None,
attributes: list[TableObjectAttribute] | None = None,
description: str | None = None,
sort_key_attribute: TableObjectAttribute | None = None,
ttl_attribute: TableObjectAttribute | None = None,
) -> type["TableObject"]:
"""
Define a TableObject
Keyword Arguments:
attribute_lookup_prefix: Attribute lookup prefix
attributes: List of attributes
description: Description of the table
partition_key_attribute: Partition key attribute
object_name: Name of the object
sort_key_attribute: Sort key attribute
table_name: Name of the table
ttl_attribute: Optional TTL attribute
"""
obj_klass = type(object_name, (cls,), {})
obj_klass.partition_key_attribute = partition_key_attribute # type: ignore[attr-defined]
obj_klass.object_name = object_name # type: ignore[attr-defined]
obj_klass.table_name = table_name # type: ignore[attr-defined]
obj_klass.attribute_lookup_prefix = attribute_lookup_prefix # type: ignore[attr-defined]
obj_klass.sort_key_attribute = sort_key_attribute # type: ignore[attr-defined]
obj_klass.attributes = attributes or [] # type: ignore[attr-defined]
obj_klass.description = description # type: ignore[attr-defined]
obj_klass.ttl_attribute = ttl_attribute # type: ignore[attr-defined]
return obj_klass
[docs]
def __getattr__(self, name: str) -> Any:
"""
Get an attribute by name
Keyword Arguments:
name -- Name of the attribute
Returns:
Any
"""
attr_keys = [attr.name for attr in self.all_attributes()]
if self.attribute_lookup_prefix:
prefixed_name = f"{self.attribute_lookup_prefix}_{name}"
if prefixed_name in attr_keys:
return getattr(self, prefixed_name, None)
return super().__getattribute__(name)
[docs]
def attribute_value(self, name: str) -> Any:
"""
Get the value of an attribute
Keyword arguments:
name -- Name of the attribute
"""
return getattr(self, name)
[docs]
def composite_attribute_values(self, name: str) -> dict:
"""
Get a dictionary representation of the composite attribute's values
Keyword arguments:
name -- Name of the attribute
"""
attr = self.attribute_definition(name)
if not attr:
raise ValueError(f"Attribute {name} not found")
if attr.attribute_type is not TableObjectAttributeType.COMPOSITE_STRING:
raise ValueError(f"Attribute {name} is not a composite string")
full_value = self.attribute_value(name)
split_values = full_value.split("-")
return {arg: split_values[idx] for idx, arg in enumerate(attr.argument_names or [])}
[docs]
def execute_on_update(self) -> None:
"""
Execute the on update function
Override this method to provide custom behavior when the object is saved to DynamoDB
"""
logging.debug("Executing default execute_on_update function ... nothing updated")
[docs]
def update(self, **kwargs) -> list[str]:
"""
Update the attributes of the object and provide a list of attribute names
that were updated.
Keyword arguments:
kwargs -- Attributes to update
"""
changed_attrs: list[str] = []
for attr in self.all_attributes():
if attr.name in kwargs:
curr_val = getattr(self, attr.name)
new_val = kwargs[attr.name]
if curr_val != new_val:
attr.set_attribute(self, new_val)
changed_attrs.append(attr.name)
return changed_attrs
[docs]
def to_dict(
self, exclude_attribute_names: list[str] | None = None, json_compatible: bool | None = False
) -> dict:
"""
Convert the object to a dict representation
Keyword Arguments:
exclude_attribute_names -- List of attribute names to exclude from the resulting dict
json_compatible -- Convert datetime objects to strings and sets to lists for JSON compatibility
"""
res: dict = {}
if exclude_attribute_names is None:
exclude_attribute_names = []
for attr in self.all_attributes():
if attr.exclude_from_dict or attr.name in exclude_attribute_names:
continue
val = getattr(self, attr.name)
if (
attr.attribute_type is TableObjectAttributeType.DATETIME
and json_compatible
and val is not None
):
val = val.isoformat()
if (
json_compatible
and attr.attribute_type is TableObjectAttributeType.STRING_SET
or attr.attribute_type is TableObjectAttributeType.NUMBER_SET
) and val is not None:
val = list(val)
if attr.custom_exporter:
val = attr.custom_exporter(val)
res[attr.name] = val
return res
[docs]
def to_dynamodb_item(self) -> dict:
"""
Convert the object to a DynamoDB item
Returns:
Dict
"""
item: dict = {}
for attr in self.all_attributes():
val = getattr(self, attr.name)
dyn_attr = attr.as_dynamodb_attribute(val)
if dyn_attr:
item.update(dyn_attr)
return item
[docs]
def to_json(self) -> str:
"""
Convert the object to a JSON string
Returns:
str
"""
return json.dumps(self.to_dict(json_compatible=True))
[docs]
@classmethod
def all_attributes(cls) -> list[TableObjectAttribute]:
"""
Class method that returns all defined attributes on the class
"""
attributes = deepcopy(cls.attributes)
attributes.append(cls.partition_key_attribute)
if cls.sort_key_attribute:
attributes.append(cls.sort_key_attribute)
if cls.ttl_attribute:
attributes.append(cls.ttl_attribute)
return attributes
[docs]
@classmethod
def attribute_definition(cls, name: str) -> TableObjectAttribute | None:
"""
Get an attribute definition by name
Keyword Arguments:
name -- Name of the attribute
Returns:
TableObjectAttribute
"""
for attr in cls.all_attributes():
if attr.name == name:
return attr
return None
[docs]
@classmethod
def from_dynamodb_item(cls, item: dict) -> "TableObject":
"""
Create a TableObject from a DynamoDB item
Keyword Arguments:
item -- Item to convert
Returns:
TableObject
"""
updated_item: dict = {}
for attr in cls.all_attributes():
if attr.dynamodb_key_name in item:
val = item[attr.dynamodb_key_name]
updated_item[attr.name] = attr.from_dynamodb_attribute(val)
return cls(**updated_item)
[docs]
@classmethod
def gen_dynamodb_key(cls, partition_key_value: str, sort_key_value: str | None = None) -> dict:
"""
Generate a DynamoDB key
Keyword Arguments:
partition_key_value -- Partition key value
sort_key_value -- Sort key value
Returns:
Dict
"""
key = cls.partition_key_attribute.as_dynamodb_attribute(partition_key_value)
if cls.sort_key_attribute:
if not sort_key_value:
raise ValueError("Sort key attribute is required, no value provided")
key.update(cls.sort_key_attribute.as_dynamodb_attribute(sort_key_value))
return key
[docs]
@classmethod
def schema_description(cls) -> str:
"""
Get the schema for the object in a human readable format
Returns:
str
"""
full_descr = cls.object_name or cls.__name__
if cls.description:
full_descr += f" - {cls.description}"
for attr in cls.all_attributes():
if attr.exclude_from_schema_description:
continue
full_descr += f"\n - {attr.schema_to_str()}"
return full_descr
[docs]
@classmethod
def schema_to_str(cls, only_indexed_attributes: bool = True) -> str:
"""
Describe the full schema for the object
Keyword Arguments:
only_indexed_attributes -- Only describe indexed attributes
Returns:
str
"""
schema_str_list = [
cls.schema_description(),
]
for attr in cls.all_attributes():
schema_str_list.append(attr.schema_to_str())
return "\n".join(schema_str_list)
[docs]
@staticmethod
def update_date_attributes(
date_attribute_names: list[str], obj: "TableObject", to_datetime: datetime | None = None
) -> None:
"""
Update the record_last_updated attribute on the object. Helper method that is
commonly used to construct execute_on_update functions.
Keyword Arguments:
date_attribute_names -- Names of the date attributes
obj -- Object to update
to_datetime -- Datetime to set, defaults to datetime.now()
"""
if not to_datetime:
to_datetime = datetime.now(tz=UTC)
for attr_name in date_attribute_names:
setattr(obj, attr_name, to_datetime)