"""Reusable serialization mixin for projected SQLAlchemy models."""
from __future__ import annotations
from collections.abc import Collection, Sequence
from datetime import date, datetime
from typing import Any, ClassVar
from edri.utility.function import camel2snake, snake2camel
from sqlalchemy import inspect as sqlalchemy_inspect
from edri.integrations.sqlalchemy.projection import ProjectionNode
__all__ = ["SQLAlchemyModelMixin"]
[docs]
class SQLAlchemyModelMixin:
"""Serialize mapped columns and relationships using EDRI projections.
Declarative bases can inherit from this mixin before ``DeclarativeBase``.
``__json__`` supplies the serialized shape only on the non-projected path
(bare ORM responses and :class:`SessionJSONEncoder`); a model used solely
through projections may leave it empty. Applications may override
:meth:`_json_extra` to add derived values or apply request-specific access
rules.
"""
__json__: ClassVar[Sequence[str]] = ()
__repr_fields__: ClassVar[Sequence[str]] = ()
def __repr__(self) -> str:
state = sqlalchemy_inspect(self)
values: list[str] = []
for name in self.__repr_fields__:
value = getattr(self, name, None) if name not in state.unloaded else "UNLOADED"
if isinstance(value, (date, datetime)):
value = value.isoformat()
elif hasattr(value, "name"):
value = value.name
values.append(f"{name}={value!r}")
return f"{type(self).__name__}({', '.join(values)})"
def _json_extra(
self,
requested_columns: list[str],
data: dict[str, Any],
session: Any = None,
skip_relationships: Collection[str] | None = None,
*,
projection: ProjectionNode | None = None,
) -> dict[str, Any]:
"""Return application-specific values appended to serialized output."""
return {}
[docs]
def to_json(
self,
columns: Sequence[str] | None = None,
session: Any = None,
skip_relationships: Collection[str] | None = None,
*,
projection: ProjectionNode | None = None,
_seen_orm: set[tuple[type[Any], tuple[Any, ...]]] | None = None,
) -> dict[str, Any]:
"""Serialize only mapped state permitted by the active projection."""
state = sqlalchemy_inspect(self)
identity = tuple(state.identity) if state.identity is not None else None
identity_key = (type(self), identity) if identity is not None else None
if _seen_orm is None:
_seen_orm = set()
elif identity_key is not None and identity_key in _seen_orm:
mapper = sqlalchemy_inspect(type(self))
return {
snake2camel(primary_key.key): getattr(self, primary_key.key)
for primary_key in mapper.primary_key
}
if identity_key is not None:
_seen_orm.add(identity_key)
try:
requested_columns = (
list(projection.visible_columns)
if projection is not None
else [camel2snake(column) for column in columns]
if columns is not None
else list(self.__json__)
)
skipped = set(skip_relationships or ())
mapper = sqlalchemy_inspect(type(self))
relationships = mapper.relationships
relationship_keys = set(relationships.keys())
mapped_keys = set(mapper.column_attrs.keys()) | relationship_keys
serializable_columns = list(requested_columns)
if projection is not None:
serializable_columns.extend(projection.visible_relationships)
data: dict[str, Any] = {}
for column in serializable_columns:
if column not in mapped_keys or column in skipped:
continue
if (
projection is not None
and column in relationship_keys
and column not in projection.visible_relationships
):
continue
value = getattr(self, column)
if column not in relationship_keys:
data[snake2camel(column)] = value
continue
relationship = relationships[column]
child_skipped = {relationship.back_populates} if relationship.back_populates else set()
child_projection = None
if projection is not None:
child_projection = (
projection
if column in projection.recursive_relationships
else projection.relationships.get(column)
)
if value is None:
data[snake2camel(column)] = None
elif relationship.uselist:
data[snake2camel(column)] = [
item.to_json(
session=session,
skip_relationships=child_skipped,
projection=child_projection,
_seen_orm=_seen_orm,
)
if hasattr(item, "to_json")
else item
for item in value
]
else:
data[snake2camel(column)] = (
value.to_json(
session=session,
skip_relationships=child_skipped,
projection=child_projection,
_seen_orm=_seen_orm,
)
if hasattr(value, "to_json")
else value
)
data.update(
self._json_extra(
requested_columns,
data,
session,
skipped,
projection=projection,
)
)
return data
finally:
if identity_key is not None:
_seen_orm.discard(identity_key)