from __future__ import annotations
from collections.abc import Mapping
from dataclasses import fields
from typing import Any
from edri.api.broker import EventInfo
from edri.api.middleware import Middleware, MiddlewareControl
from edri.dataclass.event import Event
from edri.dataclass.response import ResponseStatus
from sqlalchemy import inspect as sqlalchemy_inspect
from sqlalchemy.engine.row import Row, RowMapping
from sqlalchemy.exc import InvalidRequestError
from sqlalchemy.orm.exc import DetachedInstanceError
from sqlalchemy.orm.state import InstanceState
from edri.integrations.sqlalchemy.projection import (
IncompleteProjectionError,
PROJECTION_VALUES_ATTRIBUTE,
ProjectionNode,
get_projection_spec,
materialize_response_default,
validate_projection_graph,
)
__all__ = ["SQLAlchemyResponseMiddleware", "normalize_event_response"]
[docs]
def normalize_event_response(event: Event) -> None:
"""Convert already-loaded ORM response values into transport values.
The manager has already completed and validated projected graphs before
queue transport. This broker-side step only creates transport values.
"""
response = event.get_response()
if response is None or response.get_status() is ResponseStatus.FAILED:
return
plans = getattr(response, "_orm_projections", {})
jwt_session: Any = getattr(event, "session", None)
pending: list[tuple[str, Any, ProjectionNode, str]] = []
orm_pending: list[tuple[str, Any, str]] = []
for response_field in fields(response):
field_name = response_field.name
if field_name.startswith("_"):
continue
value = materialize_response_default(response, response_field)
spec = get_projection_spec(response, field_name)
if spec is not None:
node = plans.get(field_name) if isinstance(plans, dict) else None
if node is None:
if _contains_orm(value):
raise IncompleteProjectionError(
f'{type(response).__name__}.{field_name} has ORM projection metadata '
"but projection() was not applied to the query."
)
# Conditional GETs answered from EDRI's cache do not execute
# the handler, so no projection plan or ORM values exist.
continue
pending.append(
(field_name, value, node, f"{type(response).__name__}.{field_name}")
)
elif _contains_orm(value):
orm_pending.append(
(field_name, value, f"{type(response).__name__}.{field_name}")
)
for _, value, node, path in pending:
validate_projection_graph(value, node, path=path)
seen_orm: set[tuple[type, tuple[Any, ...]]] = set()
for field_name, value, node, _ in pending:
object.__setattr__(
response,
field_name,
_normalize(
value,
node,
jwt_session,
path=f"{type(response).__name__}.{field_name}",
active=set(),
seen_orm=seen_orm,
),
)
for field_name, value, path in orm_pending:
object.__setattr__(
response,
field_name,
_normalize_bare_orm(
value,
jwt_session,
path=path,
active=set(),
seen_orm=seen_orm,
),
)
# Drop the projection plan once it has been applied. This keeps the
# response lightweight and, more importantly, prevents the broker-side
# middleware from re-validating an already-normalized response after the
# manager has converted it to transport dicts.
if hasattr(response, "_orm_projections"):
object.__setattr__(response, "_orm_projections", {})
[docs]
class SQLAlchemyResponseMiddleware(Middleware):
"""Convert already-loaded ORM response values into transport values."""
[docs]
def process_response(
self,
event: Event,
info: EventInfo | None,
) -> MiddlewareControl:
normalize_event_response(event)
return MiddlewareControl.continue_processing()
def _contains_orm(value: Any, *, active: set[int] | None = None) -> bool:
"""Return whether a transport value contains ORM while rejecting cycles."""
if isinstance(sqlalchemy_inspect(value, raiseerr=False), InstanceState):
return True
if isinstance(value, RowMapping):
values = value.values()
elif isinstance(value, Row):
values = value._mapping.values()
elif isinstance(value, Mapping):
values = value.values()
elif isinstance(value, (list, tuple, set, frozenset)):
values = value
else:
return False
if active is None:
active = set()
if id(value) in active:
raise IncompleteProjectionError("Response contains a cyclic container.")
active.add(id(value))
try:
return any(_contains_orm(item, active=active) for item in values)
finally:
active.remove(id(value))
def _normalize(
value: Any,
node: ProjectionNode,
jwt_session: Any,
*,
path: str,
active: set[int],
seen_orm: set[tuple[type, tuple[Any, ...]]],
) -> Any:
if value is None:
return None
if isinstance(value, RowMapping):
return {
_transport_key(key): _normalize(
item, node, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value.items()
}
if isinstance(value, Row):
return {
_transport_key(key): _normalize(
item, node, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value._mapping.items()
}
if isinstance(value, Mapping):
return {
_transport_key(key): _normalize(
item, node, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value.items()
}
if isinstance(value, (list, tuple, set, frozenset)):
return [
_normalize(
item, node, jwt_session, path=f"{path}[{index}]", active=active, seen_orm=seen_orm
)
for index, item in enumerate(value)
]
state = sqlalchemy_inspect(value, raiseerr=False)
if not isinstance(state, InstanceState):
return value
if id(value) in active:
raise IncompleteProjectionError(
f"A cyclic relationship was encountered while normalizing {node.model.__name__}."
)
active.add(id(value))
try:
try:
result = value.to_json(
projection=node,
session=jwt_session,
_seen_orm=seen_orm,
)
_append_projection_functions(result, value, node, path=path)
return _deep_normalize(
result,
jwt_session,
path=path,
active=active,
seen_orm=seen_orm,
)
except (InvalidRequestError, DetachedInstanceError) as error:
raise IncompleteProjectionError(
f"{path} model hook _json_extra() accessed an unloaded or unplanned attribute."
) from error
finally:
active.remove(id(value))
def _append_projection_functions(
result: dict[str, Any],
value: Any,
node: ProjectionNode,
*,
path: str,
) -> None:
if not node.functions:
return
values = getattr(value, PROJECTION_VALUES_ATTRIBUTE, None)
if not isinstance(values, dict):
raise IncompleteProjectionError(
f"{path} is missing computed projection values. Execute the query "
"through the projection result materializer."
)
for function in node.functions:
if function.output not in values:
raise IncompleteProjectionError(
f'{path} is missing computed projection value "{function.output}".'
)
result[function.output] = values[function.output]
def _transport_key(key: Any) -> Any:
if key is None or isinstance(key, (str, int, float, bool)):
return key
return str(key)
def _normalize_bare_orm(
value: Any,
jwt_session: Any,
*,
path: str,
active: set[int],
seen_orm: set[tuple[type, tuple[Any, ...]]],
) -> Any:
"""Normalize an ORM value that has no projection plan."""
if value is None:
return None
if isinstance(value, RowMapping):
return {
_transport_key(key): _normalize_bare_orm(
item, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value.items()
}
if isinstance(value, Row):
return {
_transport_key(key): _normalize_bare_orm(
item, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value._mapping.items()
}
if isinstance(value, Mapping):
return {
_transport_key(key): _normalize_bare_orm(
item, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value.items()
}
if isinstance(value, (list, tuple, set, frozenset)):
return [
_normalize_bare_orm(
item, jwt_session, path=f"{path}[{index}]", active=active, seen_orm=seen_orm
)
for index, item in enumerate(value)
]
state = sqlalchemy_inspect(value, raiseerr=False)
if not isinstance(state, InstanceState):
return value
missing = set(getattr(value, "__json__", ())) & (
set(state.unloaded) | set(state.expired_attributes)
)
if missing:
attribute = next(iter(sorted(missing)))
raise IncompleteProjectionError(
f"{path}.{attribute} is unloaded on a bare ORM response. "
"Add orm_projection() metadata for this response field."
)
if id(value) in active:
raise IncompleteProjectionError(
f"A cyclic relationship was encountered while normalizing {type(value).__name__}."
)
active.add(id(value))
try:
try:
result = value.to_json(session=jwt_session, _seen_orm=seen_orm)
return _deep_normalize(
result,
jwt_session,
path=path,
active=active,
seen_orm=seen_orm,
)
except (InvalidRequestError, DetachedInstanceError) as error:
raise IncompleteProjectionError(
f"{path} model hook _json_extra() accessed an unloaded attribute."
) from error
finally:
active.remove(id(value))
def _deep_normalize(
value: Any,
jwt_session: Any,
*,
path: str,
active: set[int],
seen_orm: set[tuple[type, tuple[Any, ...]]],
) -> Any:
"""Recursively normalize any ORM objects that may appear in dict values."""
if value is None:
return None
if isinstance(value, RowMapping):
return {
_transport_key(key): _deep_normalize(
item, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value.items()
}
if isinstance(value, Row):
return {
_transport_key(key): _deep_normalize(
item, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value._mapping.items()
}
if isinstance(value, Mapping):
return {
_transport_key(key): _deep_normalize(
item, jwt_session, path=f"{path}[{key!r}]", active=active, seen_orm=seen_orm
)
for key, item in value.items()
}
if isinstance(value, (list, tuple, set, frozenset)):
return [
_deep_normalize(
item, jwt_session, path=f"{path}[{index}]", active=active, seen_orm=seen_orm
)
for index, item in enumerate(value)
]
state = sqlalchemy_inspect(value, raiseerr=False)
if not isinstance(state, InstanceState):
return value
if id(value) in active:
mapper = sqlalchemy_inspect(type(value))
return {
_transport_key(primary_key.key): getattr(value, primary_key.key)
for primary_key in mapper.primary_key
}
active.add(id(value))
try:
result = value.to_json(session=jwt_session, _seen_orm=seen_orm)
return _deep_normalize(
result,
jwt_session,
path=path,
active=active,
seen_orm=seen_orm,
)
finally:
active.remove(id(value))