Source code for edri.integrations.sqlalchemy.event

from __future__ import annotations

from dataclasses import dataclass
from types import UnionType
from typing import Any, ClassVar, TypeVar, Type, get_type_hints, get_origin, get_args, Union

from edri.dataclass.event import Event
from edri.integrations.sqlalchemy._core import ForeignKeyValidator
from edri.integrations.sqlalchemy.session import register_orm_session
from edri.utility.function import Injection
from sqlalchemy.orm import Session


__all__ = ["ORMEvent", "orm_event"]

T = TypeVar("T", bound="ORMEvent")




@dataclass
class ORMEvent(Event):
    # per-class: cache with information about which fields are ORM / list-of-ORM etc.
    # example: {"information": "single", "parents": "list", "role_group": "list", ...}
    _orm_fields: ClassVar[dict[str, str]] = {}

    # --------- class-level helpers (run once per subclass) ---------

    @classmethod
    def _init_orm_field_cache(cls) -> None:
        """
        Initialize _orm_fields for this subclass.
        Intended to be called from the @orm_event decorator.
        """
        if cls._orm_fields:
            return  # already initialized

        # include_extras=True so generics / future stuff survives
        hints = get_type_hints(cls, include_extras=True)

        orm_field_kinds: dict[str, str] = {}

        for name, annotation in hints.items():
            kind = cls._classify_orm_field(annotation)
            if kind is not None:
                orm_field_kinds[name] = kind

        cls._orm_fields = orm_field_kinds

    @classmethod
    def _classify_orm_field(cls, annotation: Any) -> str | None:
        """
        Decide if given annotation is
        - a single FK -> returns "single"
        - a list[FK]  -> returns "list"
        - else        -> returns None
        """
        origin = get_origin(annotation)

        # --- NEW: unwrap Optional / Union[T, None] ---
        if origin is Union or origin is UnionType:
            args = [a for a in get_args(annotation) if a is not type(None)]
            if len(args) == 1:
                # Re-run classification on the inner type
                return cls._classify_orm_field(args[0])
            return None  # multi-type unions not supported

        # Handle list[inject(ForeignKeyValidator[Model], session_factory=SessionFactory)]
        if origin is list:
            (inner,) = get_args(annotation)
            if cls._is_injection_fk(inner):
                return "list"
            return None

        # Handle plain inject(ForeignKeyValidator[Model], session_factory=SessionFactory)
        if cls._is_injection_fk(annotation):
            return "single"

        return None

    @staticmethod
    def _is_injection_fk(annotation: Any) -> bool:
        """
        True if the annotation is an Injection that contains
        a ForeignKeyValidator[...] class in its 'classes'.
        """
        if not isinstance(annotation, Injection):
            return False

        for injected_cls in annotation.classes:
            if ORMEvent._is_fk_type(injected_cls):
                return True

        return False

    @staticmethod
    def _is_fk_type(tp: Any) -> bool:
        """
        True if tp is ForeignKeyValidator[...] or a subclass thereof.
        """
        # Handles things like ForeignKeyValidator[Model]
        origin = get_origin(tp) or tp
        try:
            return issubclass(origin, ForeignKeyValidator)
        except TypeError:
            return False

    # --------- merge logic using cached info ---------



    def merge(self, session: Session, load: bool = False) -> None:
        """
        Merge all mapped ORM instances in this event into the given Session.
        Uses cached field info in _orm_fields to be fast.
        (Assumes @orm_event has already populated _orm_fields.)
        """
        # optional safety: in case someone forgot the decorator
        if not type(self)._orm_fields:
            type(self)._init_orm_field_cache()

        register_orm_session(
            session,
        )

        for field_name, kind in type(self)._orm_fields.items():
            value = getattr(self, field_name)

            if value is None:
                continue

            if kind == "single":
                merged = session.merge(value, load=load)
                setattr(self, field_name, merged)

            elif kind == "list":
                merged_list = [session.merge(obj, load=load) for obj in value]
                setattr(self, field_name, merged_list)




# --------- the decorator that populates _orm_fields ---------



def orm_event(cls: Type[T]) -> Type[T]:
    """
    Class decorator for SQLEvent subclasses.

    It precomputes the _orm_fields mapping at import time so merge()
    doesn’t need to do any reflection.
    """
    if not issubclass(cls, ORMEvent):
        raise TypeError("@orm_event can only be used on SQLEvent subclasses")

    cls._init_orm_field_cache()
    return cls