"""Manager-local SQLAlchemy response finalization.
ORM values remain ORM values until the broker middleware.
This module keeps ORM sessions alive for one ManagerBase.resolve() call,
completes projected response graphs while their original sessions are still
available, validates the graphs, and closes the registered sessions before
the event enters multiprocessing.
"""
from __future__ import annotations
from collections.abc import Mapping
from contextvars import ContextVar, Token
from dataclasses import fields
from multiprocessing.connection import Connection
from typing import Any, ClassVar, Self
from edri.abstract import ManagerBase
from edri.dataclass.event import Event
from sqlalchemy import inspect as sqlalchemy_inspect, select
from sqlalchemy.orm import Session
from sqlalchemy.orm.state import InstanceState
__all__ = [
"ORMResolutionScope",
"ResponseAwareSession",
"SQLAlchemyManager",
"SQLAlchemyManagerMixin",
"finalize_orm_response",
"register_orm_session",
]
[docs]
class ORMResolutionScope:
"""Keep ORM sessions alive for one manager resolution."""
_current: ClassVar[ContextVar[ORMResolutionScope | None]] = ContextVar(
"orm_resolution_scope",
default=None,
)
def __init__(self) -> None:
self._sessions: list[Session] = []
self._token: Token[ORMResolutionScope | None] | None = None
def __enter__(self) -> Self:
if self._token is not None:
raise RuntimeError(
"ORMResolutionScope is already active."
)
self._token = self._current.set(self)
return self
def __exit__(self, *_: Any) -> None:
try:
self.close_sessions()
finally:
token = self._token
self._token = None
if token is not None:
self._current.reset(token)
[docs]
@classmethod
def current(cls) -> ORMResolutionScope | None:
"""Return the scope active in the current execution context."""
return cls._current.get()
[docs]
def register(self, session: Session) -> None:
"""Register a session for deferred closing."""
if session not in self._sessions:
self._sessions.append(session)
[docs]
def close_sessions(self) -> None:
"""Close all registered sessions in reverse registration order.
``Session.close(session)`` deliberately bypasses
``ResponseAwareSession.close()`` so the close is performed immediately
instead of being deferred again.
"""
while self._sessions:
Session.close(
self._sessions.pop()
)
[docs]
class ResponseAwareSession(Session):
"""Session behavior required by ORM events and projected responses."""
[docs]
def merge(
self,
instance: Any,
*,
load: bool = True,
options: Any = None,
) -> Any:
suspended: list[tuple[InstanceState[Any], Any, Any]] = []
_suspend_raiseload_state(instance, seen=set(), suspended=suspended)
try:
merged = super().merge(instance, load=load, options=options)
finally:
for state, callables, load_options in reversed(suspended):
state.callables = callables
state.load_options = load_options
_reset_merged_loader_state(merged, seen=set())
return merged
[docs]
def close(self) -> None:
scope = ORMResolutionScope.current()
if scope is None:
super().close()
return
scope.register(self)
def _reset_merged_loader_state(value: Any, *, seen: set[int]) -> None:
"""Remove query-local loader restrictions copied by ``Session.merge``.
Projection queries use per-instance raiseload callables to prevent
accidental SQL while a response graph is being finalized. SQLAlchemy
copies those callables and their loader options to the destination of a
merge. The destination belongs to a new session, where normal lazy loads
are valid again, so retaining the projection restrictions is incorrect.
"""
state = sqlalchemy_inspect(value, raiseerr=False)
if isinstance(state, InstanceState):
state_id = id(state)
if state_id in seen:
return
seen.add(state_id)
raiseload_keys = {
key
for key, loader in state.callables.items()
if _is_raiseload_callable(loader)
}
if raiseload_keys:
state.expired_attributes.update(raiseload_keys.difference(state.dict))
state.callables = {
key: loader
for key, loader in state.callables.items()
if key not in raiseload_keys
}
state.load_options = ()
for child in state.dict.values():
_reset_merged_loader_state(child, seen=seen)
return
if isinstance(value, Mapping):
children = value.values()
elif isinstance(value, (list, tuple, set, frozenset)):
children = value
else:
return
for child in children:
_reset_merged_loader_state(child, seen=seen)
def _suspend_raiseload_state(
value: Any,
*,
seen: set[int],
suspended: list[tuple[InstanceState[Any], Any, Any]],
) -> None:
"""Temporarily prevent ``Session.merge`` from copying raiseload state."""
state = sqlalchemy_inspect(value, raiseerr=False)
if isinstance(state, InstanceState):
state_id = id(state)
if state_id in seen:
return
seen.add(state_id)
if any(_is_raiseload_callable(loader) for loader in state.callables.values()):
suspended.append((state, state.callables, state.load_options))
state.callables = {
key: loader
for key, loader in state.callables.items()
if not _is_raiseload_callable(loader)
}
state.load_options = ()
for child in state.dict.values():
_suspend_raiseload_state(child, seen=seen, suspended=suspended)
return
if isinstance(value, Mapping):
children = value.values()
elif isinstance(value, (list, tuple, set, frozenset)):
children = value
else:
return
for child in children:
_suspend_raiseload_state(child, seen=seen, suspended=suspended)
def _is_raiseload_callable(loader: Any) -> bool:
if getattr(loader, "raiseload", False):
return True
strategy = dict(getattr(loader, "strategy_key", ()))
return strategy.get("lazy") in {"raise", "raise_on_sql"}
[docs]
def register_orm_session(
session: Session,
) -> None:
"""Register an ORM session for the current manager resolution."""
scope = ORMResolutionScope.current()
if scope is not None:
scope.register(session)
[docs]
def finalize_orm_response(event: Event) -> None:
"""Complete and validate projected ORM response graphs.
ORM instances remain ORM instances. Serialization to dictionaries happens
later in the broker-side SQLAlchemy response middleware.
If a response field already has a resolved ProjectionNode, this function
only validates that all requested state is available.
If the ProjectionNode is missing, each ORM root must still be attached to
its original Session. Roots are grouped by Session and the projection query
is executed once per Session, using ``populate_existing=True`` so the
original response instances are populated rather than unrelated instances
from another Session.
Any attached Session encountered here is registered with the active
ORMResolutionScope and all registered Sessions are closed after successful
finalization, detaching the response graph before multiprocessing.
"""
response = event.get_response()
if response is None:
return
from edri.integrations.sqlalchemy.projection import (
get_projection_spec,
materialize_response_default,
projection,
projection_all,
validate_projection_graph,
)
scope = ORMResolutionScope.current()
plans = getattr(
response,
"_orm_projections",
{},
)
for response_field in fields(response):
field_name = response_field.name
spec = get_projection_spec(
response,
field_name,
)
if spec is None:
continue
value = materialize_response_default(
response,
response_field,
)
node = (
plans.get(field_name)
if isinstance(plans, dict)
else None
)
model = (
node.model
if node is not None
else spec.model
)
roots = _roots(
value,
model,
)
if not roots:
continue
roots_by_session = _group_and_register_roots_by_session(
roots,
scope=scope,
response=response,
field_name=field_name,
require_attached=node is None,
)
if node is None:
if model is None:
raise RuntimeError(
f"{type(response).__name__}.{field_name} "
"does not define an ORM model and the model could not "
"be inferred from a resolved projection."
)
if scope is None:
raise RuntimeError(
f"{type(response).__name__}.{field_name} needs an active "
"ORMResolutionScope for projection finalization."
)
primary_key = _primary_key(
model,
response=response,
field_name=field_name,
)
statement = projection(
select(model),
event=event,
field=field_name,
)
for session, session_roots in roots_by_session.items():
identities = _identities(
session_roots,
response=response,
field_name=field_name,
)
with session.no_autoflush:
projection_all(
session,
statement
.where(
primary_key.in_(
identities
)
)
.execution_options(
populate_existing=True,
),
response=response,
field=field_name,
)
plans = getattr(
response,
"_orm_projections",
{},
)
if not isinstance(
plans,
dict,
):
raise RuntimeError(
f"{type(response).__name__}._orm_projections "
"must be a dictionary."
)
try:
node = plans[field_name]
except KeyError as error:
raise RuntimeError(
"projection() did not store a resolved plan for "
f"{type(response).__name__}.{field_name}."
) from error
for index, root in enumerate(roots):
validate_projection_graph(
root,
node,
path=(
f"{type(response).__name__}."
f"{field_name}[{index}]"
),
)
if scope is not None:
scope.close_sessions()
def _group_and_register_roots_by_session(
roots: list[Any],
*,
scope: ORMResolutionScope | None,
response: Any,
field_name: str,
require_attached: bool,
) -> dict[Session, list[Any]]:
"""Group roots by their owning Session and register those Sessions.
An attached ORM response cannot safely leave finalization without an
active resolution scope because no component would then be responsible for
detaching it before multiprocessing. Detached roots are valid when their
projection is already resolved, but cannot be completed by another query.
"""
grouped: dict[Session, list[Any]] = {}
for index, root in enumerate(roots):
state = sqlalchemy_inspect(root)
session = state.session
if session is None:
if require_attached:
raise RuntimeError(
f"{type(response).__name__}.{field_name}[{index}] contains "
f"detached {type(root).__name__}, but the response "
"projection still requires finalization."
)
continue
if scope is None:
raise RuntimeError(
f"{type(response).__name__}.{field_name}[{index}] contains "
f"attached {type(root).__name__}, but no "
"ORMResolutionScope is active."
)
scope.register(session)
grouped.setdefault(
session,
[],
).append(root)
return grouped
def _primary_key(
model: type[Any],
*,
response: Any,
field_name: str,
) -> Any:
"""Return the model's single primary-key column."""
primary_keys = tuple(
sqlalchemy_inspect(
model
).primary_key
)
if len(primary_keys) != 1:
raise RuntimeError(
f"{type(response).__name__}.{field_name} uses "
f"{model.__name__}, which has a composite primary key. "
"The current response finalizer supports only models "
"with one primary-key column."
)
return primary_keys[0]
def _identities(
roots: list[Any],
*,
response: Any,
field_name: str,
) -> list[Any]:
"""Return unique persistent identities for ORM roots."""
identities: list[Any] = []
for index, root in enumerate(roots):
state = sqlalchemy_inspect(root)
identity = state.identity
if identity is None:
raise RuntimeError(
f"{type(response).__name__}.{field_name}[{index}] contains "
f"{type(root).__name__} without a persistent database "
"identity."
)
if len(identity) != 1:
raise RuntimeError(
f"{type(response).__name__}.{field_name}[{index}] contains "
f"{type(root).__name__} with a composite identity, which "
"is not supported by the current response finalizer."
)
identities.append(
identity[0]
)
return list(
dict.fromkeys(
identities
)
)
def _roots(
value: Any,
model: type[Any] | None,
) -> list[Any]:
"""Return immediate ORM roots from a projected response field."""
if model is None:
return []
if isinstance(
value,
Mapping,
):
values = value.values()
elif isinstance(
value,
(
list,
tuple,
set,
frozenset,
),
):
values = value
else:
return (
[value]
if isinstance(
value,
model,
)
else []
)
return [
item
for item in values
if isinstance(
item,
model,
)
]
[docs]
class SQLAlchemyManagerMixin:
"""Add SQLAlchemy response finalization to an EDRI manager.
Place the mixin before the concrete manager base so cooperative ``super()``
calls reach :class:`~edri.abstract.ManagerBase` or one of its variants.
"""
[docs]
def resolve(self, event: Event) -> None:
with ORMResolutionScope():
super().resolve(event)
[docs]
def resolve_non_router_event(
self,
event: Event,
pipe: Connection,
) -> None:
with ORMResolutionScope():
super().resolve_non_router_event(event, pipe)
[docs]
def prepare_outgoing_event(self, event: Event) -> None:
if event.has_response():
finalize_orm_response(event)
super().prepare_outgoing_event(event)
[docs]
class SQLAlchemyManager(SQLAlchemyManagerMixin, ManagerBase):
"""Ready-to-use synchronous manager with SQLAlchemy response support."""