from __future__ import annotations
from collections.abc import Mapping
from dataclasses import Field, MISSING, dataclass, fields
from types import UnionType
from typing import Any, Union, get_args, get_origin, get_type_hints
from edri.dataclass.response import Response
from sqlalchemy import func, inspect as sqlalchemy_inspect, select
from sqlalchemy.engine.row import Row, RowMapping
from sqlalchemy.orm import Load, Session, load_only, selectinload
from sqlalchemy.orm.state import InstanceState
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.sql.selectable import Select
from edri.integrations.sqlalchemy.path import SQLAlchemyPathError, canonical_orm_path
__all__ = [
"IncompleteProjectionError",
"ProjectionError",
"ProjectionFunctionSpec",
"ProjectionNode",
"ProjectionSpec",
"get_projection_spec",
"orm_count",
"orm_projection",
"projection",
"projection_all",
"projection_one",
"projection_rows",
]
PROJECTION_METADATA_KEY = "orm_projection"
MAX_RECURSIVE_LOAD_DEPTH = 64
PROJECTION_VALUES_ATTRIBUTE = "_edri_projection_values"
[docs]
class ProjectionError(ValueError):
"""Raised when response projection metadata or a requested path is invalid."""
[docs]
class IncompleteProjectionError(RuntimeError):
"""Raised when a response ORM graph does not match its resolved plan."""
[docs]
@dataclass(frozen=True, slots=True)
class ProjectionFunctionSpec:
function: str
path: str
output: str
@property
def request(self) -> str:
return f"{self.function}({self.path})"
[docs]
def orm_count(
path: str,
*,
output: str | None = None,
) -> ProjectionFunctionSpec:
"""Declare an allowed COUNT projection over a direct relationship."""
if not isinstance(path, str) or not path.strip():
raise ValueError("COUNT projection path cannot be empty")
path = path.strip()
if "." in path:
raise ValueError("COUNT currently supports direct relationships only")
request = f"count({path})"
if output is None:
output = request
if (
output != request
and (not output or not output.replace("_", "").isalnum())
):
raise ValueError(f'Invalid COUNT projection output name "{output}".')
return ProjectionFunctionSpec(function="count", path=path, output=output)
[docs]
@dataclass(frozen=True, slots=True)
class ProjectionSpec:
source: str
default: tuple[str, ...]
allowed: frozenset[str]
dependencies: tuple[str, ...] = ()
model: type[Any] | None = None
recursive: tuple[str, ...] = ()
functions: tuple[ProjectionFunctionSpec, ...] = ()
@dataclass(frozen=True, slots=True)
class ResolvedProjectionFunction:
function: str
path: tuple[str, ...]
output: str
result_key: str
[docs]
@dataclass(frozen=True, slots=True)
class ProjectionNode:
model: type[Any]
visible_columns: tuple[str, ...]
loaded_columns: tuple[str, ...]
relationships: dict[str, "ProjectionNode"]
visible_relationships: tuple[str, ...]
recursive_relationships: tuple[str, ...] = ()
functions: tuple[ResolvedProjectionFunction, ...] = ()
@dataclass
class _NodeBuilder:
model: type[Any]
visible_columns: list[str]
loaded_columns: list[str]
relationships: dict[str, "_NodeBuilder"]
visible_relationships: list[str]
recursive_relationships: list[str]
[docs]
def orm_projection(
*,
source: str,
default: tuple[str, ...] = (),
allowed: tuple[str, ...] = (),
dependencies: tuple[str, ...] = (),
model: type[Any] | None = None,
recursive: tuple[str, ...] = (),
functions: tuple[ProjectionFunctionSpec, ...] = (),
) -> dict[str, ProjectionSpec]:
"""Return metadata for an ORM-backed EDRI response field."""
if not source:
raise ValueError("Projection metadata source cannot be empty")
requests: set[str] = set()
outputs: set[str] = set()
for function in functions:
if function.request in requests:
raise ValueError(f'Duplicate projection function "{function.request}".')
if function.output in outputs:
raise ValueError(
f'Duplicate projection function output "{function.output}".'
)
requests.add(function.request)
outputs.add(function.output)
return {
PROJECTION_METADATA_KEY: ProjectionSpec(
source=source,
default=tuple(default),
allowed=frozenset(allowed),
dependencies=tuple(dependencies),
model=model,
recursive=tuple(recursive),
functions=tuple(functions),
)
}
[docs]
def get_projection_spec(response: Response, field_name: str) -> ProjectionSpec | None:
"""Read projection metadata while isolating the current EDRI workaround."""
response_fields = {response_field.name: response_field for response_field in fields(response)}
response_field = response_fields.get(field_name)
if response_field is None:
raise ProjectionError(
f'Response field "{field_name}" does not exist on {type(response).__name__}.'
)
metadata = response_field.metadata
candidate = metadata.get(PROJECTION_METADATA_KEY)
retained_field = response_field.default
if candidate is None and isinstance(retained_field, Field):
candidate = retained_field.metadata.get(PROJECTION_METADATA_KEY)
if candidate is None:
return None
if not isinstance(candidate, ProjectionSpec):
raise TypeError(
f'Invalid ORM projection metadata for {type(response).__name__}.{field_name}.'
)
return candidate
[docs]
def projection(
statement: Select[Any],
*,
event: Any,
field: str,
) -> Select[Any]:
"""Validate a response selection and apply its SQLAlchemy loader plan."""
response = getattr(event, "response", None)
if response is None:
raise ProjectionError("Projection event has no response object")
spec = get_projection_spec(response, field)
if spec is None:
raise ProjectionError(
f'{type(response).__name__}.{field} has no ORM projection metadata.'
)
model = _response_model(response, field, spec)
_ensure_model_is_selected(statement, model, field)
allowed = _canonical_paths(model, spec.allowed, field)
requested_raw = getattr(event, spec.source, None)
if requested_raw:
requested: tuple[str, ...] = tuple(requested_raw)
elif spec.default:
requested = spec.default
else:
# No explicit selection and no declared default: emit the whole
# allowlist. The projection contract is self-contained -- the model's
# ``__json__`` shape only drives the non-projected serialization path.
requested = tuple(sorted(spec.allowed))
requested_paths, selected_functions = _resolve_projection_functions(
model,
requested,
spec.functions,
field,
)
visible_paths = _canonical_requested_paths(
model,
requested_paths,
allowed,
field,
)
dependency_paths = _canonical_dependency_paths(model, spec.dependencies, field)
builder = _NodeBuilder(model, [], [], {}, [], [])
for path in visible_paths:
_add_path(builder, path, visible=True, response_field=field)
for path in dependency_paths:
_add_path(builder, path, visible=False, response_field=field)
recursive_paths = _canonical_recursive_paths(model, spec.recursive, field)
visible_path_set = set(visible_paths)
for path in recursive_paths:
if path in visible_path_set:
_mark_recursive(builder, path)
node = _freeze_node(builder, functions=selected_functions)
statement = statement.options(*_loader_options(node))
for function in selected_functions:
statement = statement.add_columns(
_projection_function_expression(model, function)
)
plans = getattr(response, "_orm_projections", None)
if not isinstance(plans, dict):
object.__setattr__(response, "_orm_projections", {})
response._orm_projections[field] = node
return statement
[docs]
def projection_rows(
session: Session,
statement: Select[Any],
*,
response: Response,
field: str,
) -> list[Row[Any]]:
"""Execute a projected query and preserve its computed values."""
node = _resolved_projection_node(response, field)
rows = list(session.execute(statement).unique().all())
if node.functions:
for row in rows:
try:
instance = row._mapping[node.model]
except KeyError as error:
raise ProjectionError(
f"Projected query does not return {node.model.__name__}."
) from error
_store_projection_values(instance, row._mapping, node)
return rows
[docs]
def projection_all(
session: Session,
statement: Select[Any],
*,
response: Response,
field: str,
) -> list[Any]:
"""Execute a projected collection query and preserve computed values."""
node = _resolved_projection_node(response, field)
if not node.functions:
return list(session.scalars(statement).unique().all())
return [
row._mapping[node.model]
for row in projection_rows(
session,
statement,
response=response,
field=field,
)
]
[docs]
def projection_one(
session: Session,
statement: Select[Any],
*,
response: Response,
field: str,
) -> Any:
"""Execute a projected single-row query and preserve computed values."""
node = _resolved_projection_node(response, field)
if not node.functions:
return session.scalars(statement).unique().one()
row = session.execute(statement).unique().one()
try:
instance = row._mapping[node.model]
except KeyError as error:
raise ProjectionError(
f"Projected query does not return {node.model.__name__}."
) from error
_store_projection_values(instance, row._mapping, node)
return instance
def materialize_response_default(
response: Response,
response_field: Field[Any],
) -> Any:
"""Materialize a dataclass field default retained by EDRI."""
value = getattr(response, response_field.name)
if not isinstance(value, Field):
return value
if value.default_factory is not MISSING:
value = value.default_factory()
elif value.default is not MISSING:
value = value.default
else:
value = None
object.__setattr__(response, response_field.name, value)
return value
def validate_projection_graph(
value: Any,
node: ProjectionNode,
*,
path: str,
active: set[int] | None = None,
) -> None:
"""Validate that an ORM response graph contains the planned state."""
if active is None:
active = set()
if value is None:
return
if isinstance(value, RowMapping):
_validate_mapping(value, node, path=path, active=active)
return
if isinstance(value, Row):
_validate_mapping(value._mapping, node, path=path, active=active)
return
if isinstance(value, Mapping):
_validate_mapping(value, node, path=path, active=active)
return
if isinstance(value, (list, tuple, set, frozenset)):
for index, item in enumerate(value):
validate_projection_graph(
item,
node,
path=f"{path}[{index}]",
active=active,
)
return
state = sqlalchemy_inspect(value, raiseerr=False)
if not isinstance(state, InstanceState):
raise IncompleteProjectionError(
f"{path} is not an instance of the planned ORM model "
f"{node.model.__name__}. Apply projection() to the query."
)
if not isinstance(value, node.model):
raise IncompleteProjectionError(
f"{path} contains {type(value).__name__}, "
f"expected {node.model.__name__}."
)
if id(value) in active:
return
missing_columns = set(node.loaded_columns) & (
set(state.unloaded) | set(state.expired_attributes)
)
if missing_columns:
raise _not_loaded(path, next(iter(sorted(missing_columns))))
_validate_projection_values(value, node, path=path)
active.add(id(value))
try:
for relationship_name, child in node.relationships.items():
if (
relationship_name in state.unloaded
or relationship_name in state.expired_attributes
):
raise _not_loaded(path, relationship_name)
validate_projection_graph(
getattr(value, relationship_name),
child,
path=f"{path}.{relationship_name}",
active=active,
)
for relationship_name in node.recursive_relationships:
if (
relationship_name in state.unloaded
or relationship_name in state.expired_attributes
):
raise _not_loaded(path, relationship_name)
validate_projection_graph(
getattr(value, relationship_name),
node,
path=f"{path}.{relationship_name}",
active=active,
)
finally:
active.remove(id(value))
def _validate_mapping(
value: Mapping[Any, Any],
node: ProjectionNode,
*,
path: str,
active: set[int],
) -> None:
for key, item in value.items():
if _contains_orm_or_container(item):
validate_projection_graph(
item,
node,
path=f"{path}[{key!r}]",
active=active,
)
def _contains_orm_or_container(value: Any) -> bool:
if isinstance(
value,
(Mapping, Row, RowMapping, list, tuple, set, frozenset),
):
return True
return isinstance(
sqlalchemy_inspect(value, raiseerr=False),
InstanceState,
)
def _not_loaded(path: str, attribute: str) -> IncompleteProjectionError:
return IncompleteProjectionError(
f"{path}.{attribute} was not loaded. Apply projection() to the query "
f'or remove "{attribute}" from the requested projection.'
)
def _response_model(response: Response, field_name: str, spec: ProjectionSpec) -> type[Any]:
if spec.model is not None:
return spec.model
try:
annotation = get_type_hints(type(response), include_extras=True)[field_name]
except (KeyError, TypeError, NameError) as error:
raise ProjectionError(
f'Cannot infer ORM model for response field "{field_name}".'
) from error
model = _unwrap_model_annotation(annotation)
if model is None:
raise ProjectionError(
f'Cannot infer ORM model for response field "{field_name}".'
)
return model
def _unwrap_model_annotation(annotation: Any) -> type[Any] | None:
origin = get_origin(annotation)
if origin in (list, tuple, set, frozenset):
args = get_args(annotation)
if not args:
return None
return _unwrap_model_annotation(args[0])
if origin in (Union, UnionType):
for argument in get_args(annotation):
if argument is not type(None):
model = _unwrap_model_annotation(argument)
if model is not None:
return model
return None
if isinstance(annotation, type) and sqlalchemy_inspect(annotation, raiseerr=False) is not None:
return annotation
return None
def _ensure_model_is_selected(statement: Select[Any], model: type[Any], field: str) -> None:
entities = {
description.get("entity")
for description in statement.column_descriptions
if description.get("entity") is not None
}
if model not in entities:
raise ProjectionError(
f'Query for response field "{field}" does not select {model.__name__}.'
)
def _canonical_paths(
model: type[Any],
paths: tuple[str, ...] | frozenset[str],
field: str,
) -> set[tuple[str, ...]]:
canonical: set[tuple[str, ...]] = set()
for path in paths:
canonical.add(_canonical_path(model, path, field))
return canonical
def _canonical_requested_paths(
model: type[Any],
requested: tuple[str, ...],
allowed: set[tuple[str, ...]],
field: str,
) -> list[tuple[str, ...]]:
result: list[tuple[str, ...]] = []
seen: set[tuple[str, ...]] = set()
for public_path in requested:
if not isinstance(public_path, str) or not public_path.strip():
raise ProjectionError(
f'Response field "{field}" contains an empty projection path.'
)
try:
path = _canonical_path(model, public_path, field)
except ProjectionError as error:
raise ProjectionError(
f'Projection path "{public_path}" is not allowed for response '
f'field "{field}".'
) from error
if path in seen:
raise ProjectionError(
f'Response field "{field}" contains duplicate projection path "{public_path}".'
)
if path not in allowed:
raise ProjectionError(
f'Projection path "{public_path}" is not allowed for response '
f'field "{field}".'
)
seen.add(path)
result.append(path)
return result
def _resolve_projection_functions(
model: type[Any],
requested: tuple[str, ...],
functions: tuple[ProjectionFunctionSpec, ...],
field: str,
) -> tuple[tuple[str, ...], tuple[ResolvedProjectionFunction, ...]]:
"""Split ordinary projection paths from explicitly allowed functions."""
allowed_functions = {function.request: function for function in functions}
ordinary: list[str] = []
selected: list[ResolvedProjectionFunction] = []
seen_functions: set[str] = set()
for item in requested:
function = allowed_functions.get(item) if isinstance(item, str) else None
if function is not None:
if item in seen_functions:
raise ProjectionError(
f'Response field "{field}" contains duplicate '
f'projection function "{item}".'
)
path = _canonical_path(model, function.path, field)
if len(path) != 1:
raise ProjectionError(
f'Projection function "{item}" for response field '
f'"{field}" must target a direct relationship.'
)
relationship = sqlalchemy_inspect(model).relationships.get(path[0])
if relationship is None:
raise ProjectionError(
f'Projection function "{item}" for response field '
f'"{field}" must target a relationship.'
)
if not relationship.uselist:
raise ProjectionError(
f'Projection function "{item}" for response field '
f'"{field}" requires a collection relationship.'
)
if relationship.mapper.class_ is model:
raise ProjectionError(
f'Projection function "{item}" for response field '
f'"{field}" does not support self-referential relationships yet.'
)
selected.append(
ResolvedProjectionFunction(
function=function.function,
path=path,
output=function.output,
result_key=f"__edri_projection_function_{len(selected)}",
)
)
seen_functions.add(item)
continue
if isinstance(item, str) and ("(" in item or ")" in item):
raise ProjectionError(
f'Projection function "{item}" is not allowed for response '
f'field "{field}".'
)
ordinary.append(item)
return tuple(ordinary), tuple(selected)
def _projection_function_expression(
model: type[Any],
function: ResolvedProjectionFunction,
) -> ColumnElement[Any]:
if function.function == "count":
return _count_relationship_expression(model, function)
raise ProjectionError(f'Unsupported projection function "{function.function}".')
def _count_relationship_expression(
model: type[Any],
function: ResolvedProjectionFunction,
) -> ColumnElement[int]:
relationship_name = function.path[0]
relationship = sqlalchemy_inspect(model).relationships.get(relationship_name)
if relationship is None:
raise ProjectionError(
f'COUNT target "{relationship_name}" is not a relationship on {model.__name__}.'
)
if not relationship.uselist:
raise ProjectionError(
f'COUNT target "{relationship_name}" is not a collection relationship '
f"on {model.__name__}."
)
target = relationship.mapper.persist_selectable
from_clause = (
target
if relationship.secondary is None
else relationship.secondary.join(target, relationship.secondaryjoin)
)
return (
select(func.count())
.select_from(from_clause)
.where(relationship.primaryjoin)
.correlate(model)
.scalar_subquery()
.label(function.result_key)
)
def _resolved_projection_node(response: Response, field: str) -> ProjectionNode:
plans = getattr(response, "_orm_projections", None)
node = plans.get(field) if isinstance(plans, dict) else None
if not isinstance(node, ProjectionNode):
raise ProjectionError(
f'{type(response).__name__}.{field} has no resolved projection plan.'
)
return node
def _store_projection_values(
instance: Any,
mapping: Mapping[Any, Any],
node: ProjectionNode,
) -> None:
current = getattr(instance, PROJECTION_VALUES_ATTRIBUTE, None)
if not isinstance(current, dict):
current = {}
setattr(instance, PROJECTION_VALUES_ATTRIBUTE, current)
for function in node.functions:
current[function.output] = mapping[function.result_key]
def _validate_projection_values(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}".'
)
def _canonical_dependency_paths(
model: type[Any],
dependencies: tuple[str, ...],
field: str,
) -> list[tuple[str, ...]]:
result: list[tuple[str, ...]] = []
seen: set[tuple[str, ...]] = set()
for public_path in dependencies:
path = _canonical_path(model, public_path, field)
if path not in seen:
seen.add(path)
result.append(path)
return result
def _canonical_path(model: type[Any], public_path: str, response_field: str) -> tuple[str, ...]:
try:
return canonical_orm_path(model, public_path)
except (SQLAlchemyPathError, TypeError):
raise ProjectionError(
f'Unsupported projection path "{public_path}" for response field "{response_field}".'
) from None
def _add_path(
builder: _NodeBuilder,
path: tuple[str, ...],
*,
visible: bool,
response_field: str,
) -> None:
mapper = sqlalchemy_inspect(builder.model)
head = path[0]
if head in mapper.column_attrs:
if len(path) != 1:
raise ProjectionError(
f'Unsupported projection path "{".".join(path)}" for response field "{response_field}".'
)
_append_unique(builder.loaded_columns, head)
if visible:
_append_unique(builder.visible_columns, head)
return
relationship = mapper.relationships.get(head)
if relationship is None:
raise ProjectionError(
f'Unsupported projection path "{".".join(path)}" for response field "{response_field}".'
)
child = builder.relationships.get(head)
if child is None:
child = _NodeBuilder(relationship.mapper.class_, [], [], {}, [], [])
builder.relationships[head] = child
if visible:
_append_unique(builder.visible_relationships, head)
if len(path) > 1:
_add_path(child, path[1:], visible=visible, response_field=response_field)
def _canonical_recursive_paths(
model: type[Any],
paths: tuple[str, ...],
field: str,
) -> list[tuple[str, ...]]:
result: list[tuple[str, ...]] = []
seen: set[tuple[str, ...]] = set()
for public_path in paths:
path = _canonical_path(model, public_path, field)
current_model = model
for index, relationship_name in enumerate(path):
mapper = sqlalchemy_inspect(current_model)
relationship = mapper.relationships.get(relationship_name)
if relationship is None:
raise ProjectionError(
f'Recursive projection path "{public_path}" for response field '
f'"{field}" must end in a relationship.'
)
target_model = relationship.mapper.class_
if index == len(path) - 1 and target_model is not current_model:
raise ProjectionError(
f'Recursive projection path "{public_path}" for response field '
f'"{field}" must end in a self-referential relationship.'
)
current_model = target_model
if path not in seen:
seen.add(path)
result.append(path)
return result
def _mark_recursive(
builder: _NodeBuilder,
path: tuple[str, ...],
) -> None:
"""Mark an explicitly selected self-reference as recursively repeatable.
The recursive child inherits the scalar projection of the node directly
preceding the self-referential relationship.
Relationships are deliberately not copied. Only visible and loaded scalar
columns are repeated through the recursive hierarchy.
"""
relationship_name = path[0]
child = builder.relationships.get(relationship_name)
if child is None:
# Recursive declarations constrain loading; they do not make an
# otherwise unrequested relationship visible or load it implicitly.
return
if len(path) > 1:
_mark_recursive(child, path[1:])
return
for column_name in builder.visible_columns:
_append_unique(child.visible_columns, column_name)
for column_name in builder.loaded_columns:
_append_unique(child.loaded_columns, column_name)
_append_unique(child.recursive_relationships, relationship_name)
if relationship_name in builder.visible_relationships:
_append_unique(child.visible_relationships, relationship_name)
def _freeze_node(
builder: _NodeBuilder,
*,
nested: bool = False,
functions: tuple[ResolvedProjectionFunction, ...] = (),
) -> ProjectionNode:
mapper = sqlalchemy_inspect(builder.model)
visible = list(builder.visible_columns)
loaded = list(builder.loaded_columns)
for primary_key in mapper.primary_key:
_append_unique(loaded, primary_key.key)
# A related resource without its identity is not useful to API clients.
# Root identities remain opt-in; every nested projection always exposes
# the child's primary key alongside the requested fields.
if nested:
_append_unique(visible, primary_key.key)
for relationship_name in builder.relationships:
relationship = mapper.relationships[relationship_name]
for local_column in relationship.local_columns:
if local_column.key in mapper.column_attrs:
_append_unique(loaded, local_column.key)
return ProjectionNode(
model=builder.model,
visible_columns=tuple(visible),
loaded_columns=tuple(loaded),
relationships={
name: _freeze_node(
child,
nested=name in builder.visible_relationships,
)
for name, child in builder.relationships.items()
},
visible_relationships=tuple(builder.visible_relationships),
recursive_relationships=tuple(builder.recursive_relationships),
functions=functions,
)
def _loader_options(node: ProjectionNode) -> list[Load]:
options = [
load_only(
*(getattr(node.model, name) for name in node.loaded_columns),
raiseload=True,
),
Load(node.model).raiseload("*"),
]
options.extend(
_relationship_loader(
getattr(node.model, relationship_name),
child,
recursive=relationship_name in child.recursive_relationships,
)
for relationship_name, child in node.relationships.items()
)
return options
def _relationship_loader(
attribute: Any,
node: ProjectionNode,
*,
recursive: bool = False,
) -> Load:
nested = [
_relationship_loader(
getattr(node.model, name),
child,
recursive=name in child.recursive_relationships,
)
for name, child in node.relationships.items()
]
loader = selectinload(
attribute,
recursion_depth=MAX_RECURSIVE_LOAD_DEPTH if recursive else None,
).load_only(
*(getattr(node.model, name) for name in node.loaded_columns),
raiseload=True,
)
if nested:
loader = loader.options(*nested)
return loader.raiseload("*")
def _append_unique(values: list[str], value: str) -> None:
if value not in values:
values.append(value)