arcanus.materia.sqlalchemy

The SQLAlchemy backend: the materia, the session wrappers, and re-exported loading options.

SqlalchemyMateria

arcanus.materia.sqlalchemy.base.SqlalchemyMateria

Bases: BaseMateria

Source code in arcanus/materia/sqlalchemy/base.py

class SqlalchemyMateria(BaseMateria):
    def __init__(self) -> None:
        super().__init__()
        self.column_type = SqlalchemyColumn
        self.expression_compiler = SqlalchemyExpressionCompiler()

    def bless(self, materia: type[Any]):
        def decorator(transmuter_cls: type[T]) -> type[T]:
            if transmuter_cls in self.formulars:
                raise RuntimeError(
                    f"Transmuter {transmuter_cls.__name__} is already blessed with {self} in {self.__class__.__name__}"
                )
            # Check if materia implements TransmuterProxied by verifying required attributes
            if not hasattr(materia, "transmuter_proxy"):
                raise TypeError(
                    f"{self.__class__.__name__} require materia must implement TransmuterProxied."
                )

            self.formulars[transmuter_cls] = materia

            # Inject __clause_element__ methods to make transmuter_cls compatible with SQLAlchemy SQL constructions
            def __clause_element__(cls: type[Transmuter]):
                return inspect(cls.__transmuter_provider__)

            setattr(
                transmuter_cls,
                "__clause_element__",
                classmethod(__clause_element__),
            )

            return transmuter_cls

        return decorator

    @staticmethod
    def transmuter_before_validator(
        transmuter_type: type[T], materia: object, info: ValidationInfo
    ) -> Any:
        inspector = inspect(cast(Inspectable[InstanceState[Any]], materia))

        # Get all loaded attributes from sqlalchemy orm instance
        # relationships/associations are excluded here to:
        # 1. prevent firing loading procedures to respect lazy loading
        # 2. avoid circular validation
        # related objects will be load and validated when they are visited
        data = {}
        association_proxies = extract_association_proxies(cast(Hashable, type(materia)))
        loaded_data = inspector.dict
        for field_name, field_info in transmuter_type.__pydantic_fields__.items():
            used_name = field_info.alias or field_name
            if used_name in loaded_data:
                if field_name in transmuter_type.model_associations:
                    data[used_name] = DefferedAssociation
                else:
                    data[used_name] = loaded_data[used_name]
            else:
                proxied = association_proxies.get(used_name)
                if proxied and proxied in loaded_data:
                    if field_name in transmuter_type.model_associations:
                        data[used_name] = DefferedAssociation
                    else:
                        data[used_name] = getattr(materia, used_name)

        # don't use a dict to hold loaded data here
        # to avoid pydantic's handler call this formulate function again and go to the else block
        # use an object instead to keep the behavior same with pydantic's original model_validate
        # with from_attributes=True which will skip the instance __init__.
        loaded = LoadedData()
        loaded.__dict__ = data

        return loaded

    @staticmethod
    def transmuter_before_construct(transmuter_type: type[T], materia: object) -> dict:
        # inspector: InstanceState = inspect(materia)

        # Get all loaded attributes from sqlalchemy orm instance
        # relationships/associations are excluded here to:
        # 1. prevent firing loading procedures to respect lazy loading
        # 2. avoid circular validation
        # related objects will be load and validated when they are visited
        data = {}
        association_proxies = extract_association_proxies(cast(Hashable, type(materia)))
        inspector = inspect(cast(Inspectable[InstanceState[Any]], materia))
        loaded_data = inspector.dict
        for field_name, field_info in transmuter_type.__pydantic_fields__.items():
            if field_name in transmuter_type.model_associations:
                continue
            used_name = field_info.alias or field_name

            if used_name in loaded_data:
                # TODO: support deferred columns?
                data[used_name] = getattr(materia, used_name)
            else:
                proxied = association_proxies.get(used_name)
                if proxied and proxied in loaded_data:
                    data[used_name] = getattr(materia, used_name)

        return data

    def load_association(self, association: Association):
        if not association.__instance__:
            raise RuntimeError(
                f"The relation '{association.field_name}' is not yet prepared with an owner instance."
            )
        try:
            orm_instance = association.__instance__.__transmuter_provided__
            used_name = association.used_name

            # When a RelationGroupMap field maps to an association_proxy
            # (e.g. proxy over a attribute_keyed_list_dict relationship),
            # SQLAlchemy's _AssociationDict cannot iterate dict-of-lists
            # values correctly.  Bypass the proxy and resolve manually:
            # return a plain dict[K, list[target_ORM]] that the association
            # can consume directly.
            if isinstance(association, RelationGroupMap):
                orm_class = type(orm_instance)
                proxies = extract_association_proxies(cast(Hashable, orm_class))
                if used_name in proxies:
                    target_rel = proxies[used_name]
                    association.__backing_attr__ = target_rel

                    mapper = inspect(cast(Inspectable[Mapper[Any]], orm_class))
                    proxy_descriptor: AssociationProxy
                    proxy_descriptor = mapper.all_orm_descriptors[used_name]  # type: ignore
                    value_attr = proxy_descriptor.value_attr

                    underlying = getattr(orm_instance, target_rel)
                    return {
                        key: [getattr(item, value_attr) for item in items]
                        for key, items in underlying.items()
                    }

            return getattr(orm_instance, used_name)
        except MissingGreenlet as missing_greenlet_error:
            association.__loaded__ = False
            raise MissingGreenlet(
                f"""Failed to load relation '{association.field_name}' of {association.__instance__.__class__.__name__} for a greenlet is expected. Are you trying to get the relation in a sync context ? Await the {association.__instance__.__class__.__name__}.{association.field_name} instance to trigger the sqlalchemy async IO first."""
            ) from missing_greenlet_error
        except InvalidRequestError as invalid_request_error:
            association.__loaded__ = False
            raise InvalidRequestError(
                f"""Failed to load relation '{association.field_name}' of {association.__instance__.__class__.__name__} for the relation's loading strategy is set to 'raise' in sqlalchemy. Specify the relationship with selectinload in statement options or change the loading strategy to 'select' or 'selectin' instead."""
            ) from invalid_request_error

    def aload_association(self, association: Association) -> Any:
        return greenlet_spawn(self.load_association, association)

    def association_loaded(self, association: Association) -> bool:
        owner = association.__instance_provider__
        if owner is None:
            return True
        state = inspect(cast(Inspectable[InstanceState[Any]], owner))
        return association.used_name not in state.unloaded

Session

arcanus.materia.sqlalchemy.database.Session

Bases: Session

Source code in arcanus/materia/sqlalchemy/database.py

class Session(SqlalchemySession):
    _validation_context: ValidationContextT
    _validation_context_manager: ValidateContextGeneratorT | None

    def __init__(
        self,
        bind: Optional[_SessionBind] = None,
        *,
        autoflush: bool = True,
        future: Literal[True] = True,
        expire_on_commit: bool = True,
        autobegin: bool = True,
        twophase: bool = False,
        binds: Optional[dict[_SessionBindKey, _SessionBind]] = None,
        enable_baked_queries: bool = True,
        info: Optional[_InfoType] = None,
        query_cls: Optional[type[Query[Any]]] = None,
        autocommit: Literal[False] = False,
        join_transaction_mode: JoinTransactionMode = "conditional_savepoint",
        close_resets_only: Union[bool, _NoArg] = _NoArg.NO_ARG,
    ) -> None:
        super().__init__(
            bind,
            autoflush=autoflush,
            future=future,
            expire_on_commit=expire_on_commit,
            autobegin=autobegin,
            twophase=twophase,
            binds=binds,
            enable_baked_queries=enable_baked_queries,
            info=info,
            query_cls=query_cls,
            autocommit=autocommit,
            join_transaction_mode=join_transaction_mode,
            close_resets_only=close_resets_only,
        )
        self._validation_context = WeakKeyDictionary()
        self._validation_context_manager = None

    def __enter__(self):
        self._validation_context_manager = validation_context(self._validation_context)
        self._validation_context_manager.__enter__()
        return super().__enter__()

    def __exit__(self, exc_type, exc_value, traceback) -> Optional[bool]:
        if self._validation_context_manager is not None:
            self._validation_context_manager.__exit__(exc_type, exc_value, traceback)
        return super().__exit__(exc_type, exc_value, traceback)

    def __iter__(self) -> Iterator[object]:
        if not self._validation_context:
            raise RuntimeError(
                "Active validation context is requried, please use a context manager 'with Session() as session' to create a session context."
            )
        for instance in super().__iter__():
            if instance in self._validation_context:
                yield self._validation_context[instance]
            if isinstance(instance, TransmuterProxied) and (
                transmuter := instance.transmuter_proxy
            ):
                yield transmuter
            yield instance  # type: ignore[reportUnreachable]

    @overload
    def execute(
        self,
        statement: TypedReturnsRows[_T],
        params: Optional[_CoreAnyExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> Result[_T]: ...

    @overload
    def execute(
        self,
        statement: UpdateBase,
        params: Optional[_CoreAnyExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> CursorResult[Any]: ...
    @overload
    def execute(
        self,
        statement: Executable,
        params: Optional[_CoreAnyExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> Result[Any]: ...
    def execute(
        self,
        statement: Executable,
        params: Optional[_CoreAnyExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        _parent_execute_state: Optional[Any] = None,
        _add_event: Optional[Any] = None,
    ) -> Result[Any]:
        result = super().execute(
            statement,
            params,
            execution_options=execution_options,
            bind_arguments=bind_arguments,
            _parent_execute_state=_parent_execute_state,
            _add_event=_add_event,
        )

        if execution_options.get("sa_top_level_orm_context", False):
            return result

        if execution_options.get("_sa_orm_load_options", {}):
            return result

        entities = resolve_statement_entities(statement)
        if entities and any(contains_transmuter_type(entity) for entity in entities):
            return AdaptedResult(
                real_result=result,
                entities=tuple(entities),
                # An UPDATE ... RETURNING refreshes the ORM row but hands back the
                # cached (stale) transmuter — force a re-sync from the fresh row.
                force_revalidate=isinstance(statement, Update)
                and bool(statement._returning),
            )  # pyright: ignore[reportReturnType]

        return result

    @overload
    def scalar(
        self,
        statement: TypedReturnsRows[tuple[_T]],
        params: Optional[_CoreSingleExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        **kw: Any,
    ) -> Optional[_T]: ...
    @overload
    def scalar(
        self,
        statement: Executable,
        params: Optional[_CoreSingleExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        **kw: Any,
    ) -> Any: ...
    def scalar(
        self,
        statement: Executable,
        params: Optional[_CoreSingleExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        **kw: Any,
    ) -> Any:
        return self.execute(
            statement=statement,
            params=params,
            execution_options=execution_options,
            bind_arguments=bind_arguments,
            **kw,
        ).scalar()

    @overload
    def scalars(
        self,
        statement: TypedReturnsRows[tuple[_T]],
        params: Optional[_CoreAnyExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        **kw: Any,
    ) -> ScalarResult[_T]: ...
    @overload
    def scalars(
        self,
        statement: Executable,
        params: Optional[_CoreAnyExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        **kw: Any,
    ) -> ScalarResult[Any]: ...
    def scalars(
        self,
        statement: Executable,
        params: Optional[_CoreAnyExecuteParams] = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
        **kw: Any,
    ) -> ScalarResult[Any]:
        return self.execute(
            statement=statement,
            params=params,
            execution_options=execution_options,
            bind_arguments=bind_arguments,
            **kw,
        ).scalars()

    def expunge(self, instance: object) -> None:
        if isinstance(instance, Transmuter):
            if instance.__transmuter_provided__ in self._validation_context:
                del self._validation_context[instance.__transmuter_provided__]
            super().expunge(instance.__transmuter_provided__)
        else:
            super().expunge(instance)

    def expunge_all(self) -> None:
        self._validation_context.clear()
        return super().expunge_all()

    def add(self, instance: object, _warn: bool = True) -> None:
        if isinstance(instance, Transmuter):
            self._validation_context[instance.__transmuter_provided__] = instance
            super().add(instance.__transmuter_provided__, _warn=_warn)
        else:
            super().add(instance, _warn=_warn)

    def _save_or_update_state(
        self,
        state: InstanceState,
    ) -> None:
        state._orphaned_outside_of_session = False
        self._save_or_update_impl(state)

        mapper = _state_mapper(state)
        for o, m, st_, dct_ in mapper.cascade_iterator(
            "save-update", state, halt_on=self._contains_state
        ):
            if isinstance(o, TransmuterProxied) and (transmuter := o.transmuter_proxy):
                self._validation_context[o] = transmuter
            self._save_or_update_impl(st_)

    def refresh(
        self,
        instance: object,
        attribute_names: Iterable[str] | None = None,
        with_for_update: ForUpdateArg | None | bool | dict[str, Any] = None,
    ) -> None:
        if isinstance(instance, Transmuter):
            if instance.__transmuter_provided__ not in self._validation_context:
                self._validation_context[instance.__transmuter_provided__] = instance
            super().refresh(instance, attribute_names, with_for_update)
            instance.revalidate()
        else:
            super().refresh(instance, attribute_names, with_for_update)

    def rollback(self) -> None:
        super().rollback()
        self._validation_context.clear()

    def merge(
        self,
        instance: T,
        *,
        load: bool = True,
        options: Sequence[ORMOption] | None = None,
    ) -> T:
        if isinstance(instance, Transmuter):
            if self._warn_on_events:
                self._flush_warning("Session.merge()")

            _recursive: dict[InstanceState[Any], object] = {}
            _resolve_conflict_map: dict[_IdentityKeyType[Any], object] = {}

            if load:
                # flush current contents if we expect to load data
                self._autoflush()

            object_mapper(instance)  # verify mapped
            autoflush = self.autoflush
            try:
                self.autoflush = False
                merged = self._merge(
                    attributes.instance_state(instance),
                    attributes.instance_dict(instance.__transmuter_provided__),
                    load=load,
                    options=options,
                    _recursive=_recursive,
                    _resolve_conflict_map=_resolve_conflict_map,
                )
                instance = type(instance).model_validate(merged)
                instance.revalidate()
                return instance
            finally:
                self.autoflush = autoflush
        else:
            return super().merge(
                instance,
                load=load,
                options=options,
            )

    def enable_relationship_loading(self, obj: BaseTransmuter) -> None:
        super().enable_relationship_loading(obj.__transmuter_provided__)
        self._validation_context[obj.__transmuter_provided__] = obj

    @overload
    def get(
        self,
        entity: type[T],
        ident: _PKIdentityArgument,
        *,
        options: Sequence[ORMOption] | None = None,
        populate_existing: bool = False,
        with_for_update: ForUpdateParameter = None,
        identity_token: Optional[Any] = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
    ) -> Optional[T]: ...
    @overload
    def get(
        self,
        entity: _EntityBindKey[_O],
        ident: _PKIdentityArgument,
        *,
        options: Sequence[ORMOption] | None = None,
        populate_existing: bool = False,
        with_for_update: ForUpdateParameter = None,
        identity_token: Optional[Any] = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
    ) -> Optional[_O]: ...
    def get(
        self,
        entity: type[T] | _EntityBindKey[_O],
        ident: _PKIdentityArgument,
        *,
        options: Sequence[ORMOption] | None = None,
        populate_existing: bool = False,
        with_for_update: ForUpdateParameter = None,
        identity_token: Optional[Any] = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
    ) -> Optional[T] | Optional[_O]:
        if isinstance(entity, type) and issubclass(entity, Transmuter):
            instance = super().get(
                # sqlalchemy materia requires transumter to have a provider blessed
                active_materia.get()[entity],  # pyright: ignore[reportArgumentType]
                ident,
                options=options,
                populate_existing=populate_existing,
                with_for_update=with_for_update,
                identity_token=identity_token,
                execution_options=execution_options,
                bind_arguments=bind_arguments,
            )
            if not instance:
                return None
            return entity.model_validate(instance)
        else:
            instance = super().get(
                entity,
                ident,
                options=options,
                populate_existing=populate_existing,
                with_for_update=with_for_update,
                identity_token=identity_token,
                execution_options=execution_options,
                bind_arguments=bind_arguments,
            )
            if isinstance(instance, Transmuter):
                return instance.__transmuter_provided__  # pyright: ignore[reportReturnType]
            return instance

    @overload
    def get_one(
        self,
        entity: type[T],
        ident: _PKIdentityArgument,
        *,
        options: Optional[Sequence[ORMOption]] = None,
        populate_existing: bool = False,
        with_for_update: ForUpdateParameter = None,
        identity_token: Optional[Any] = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
    ) -> T: ...
    @overload
    def get_one(
        self,
        entity: _EntityBindKey[_O],
        ident: _PKIdentityArgument,
        *,
        options: Optional[Sequence[ORMOption]] = None,
        populate_existing: bool = False,
        with_for_update: ForUpdateParameter = None,
        identity_token: Optional[Any] = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
    ) -> _O: ...
    def get_one(
        self,
        entity: type[T] | _EntityBindKey[_O],
        ident: _PKIdentityArgument,
        *,
        options: Optional[Sequence[ORMOption]] = None,
        populate_existing: bool = False,
        with_for_update: ForUpdateParameter = None,
        identity_token: Optional[Any] = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: Optional[_BindArguments] = None,
    ) -> T | _O:
        instance = self.get(
            entity,
            ident,
            options=options,
            populate_existing=populate_existing,
            with_for_update=with_for_update,
            identity_token=identity_token,
            execution_options=execution_options,
            bind_arguments=bind_arguments,
        )

        if instance is None:
            raise exc.NoResultFound("No row was found when one was required")

        return instance

    def one(
        self,
        entity: type[_T],
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool]] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ):
        statement = select(entity)

        if expressions:
            statement = statement.where(*expressions)
        if filters:
            statement = statement.filter_by(**filters)
        if options:
            statement = statement.options(*options)
        if execution_options:
            statement = statement.execution_options(**execution_options)

        return self.execute(statement).scalar_one()

    def one_or_none(
        self,
        entity: type[_T],
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool]] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ):
        statement = select(entity)

        if expressions:
            statement = statement.where(*expressions)
        if filters:
            statement = statement.filter_by(**filters)
        if options:
            statement = statement.options(*options)
        if execution_options:
            statement = statement.execution_options(**execution_options)

        return self.execute(statement).scalar_one_or_none()

    def first(
        self,
        entity: type[_T],
        order_bys: Iterable[_ColumnExpressionOrStrLabelArgument[Any]] | None = None,
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool]] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ):
        statement = select(entity)

        if order_bys:
            statement = statement.order_by(*order_bys)
        if expressions:
            statement = statement.where(*expressions)
        if filters:
            statement = statement.filter_by(**filters)
        if options:
            statement = statement.options(*options)
        if execution_options:
            statement = statement.execution_options(**execution_options)

        return self.execute(statement).scalars().first()

    def bulk(
        self,
        entity: type[_T],
        idents: Sequence[_PKIdentityArgument],
        *,
        options: Sequence[ORMOption] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
    ) -> list[_T | None]:
        """Bulk version of Session.get. Each element in idents should be
        exactly the same format as Session.get's ident parameter.

        Returns a list of entities in the same order as idents, with None
        for any ident that was not found.
        """
        if not idents:
            return []

        mapper = inspect(active_materia.get()[entity])
        pk_columns = mapper.primary_key  # pyright: ignore[reportOptionalMemberAccess]

        if len(pk_columns) == 1:
            # Build the WHERE clause based on single or composite PK
            pk_col = pk_columns[0]
            statement = select(entity).where(pk_col.in_(idents))
        else:
            # Composite PK: use tuple comparison
            # Each ident should be a tuple matching the PK columns
            statement = select(entity).where(tuple_(*pk_columns).in_(idents))

        if options:
            statement = statement.options(*options)
        if execution_options:
            statement = statement.execution_options(**execution_options)

        entities = self.execute(statement).scalars().all()

        # Build mapping from PK value(s) to entity
        if len(pk_columns) == 1:
            pk_attr = pk_columns[0].key
            mapping = {getattr(e, pk_attr): e for e in entities}
            return [mapping.get(ident) for ident in idents]
        else:
            # Composite PK: map tuple of PK values to entity
            pk_attrs = [col.key for col in pk_columns]
            mapping = {
                tuple(getattr(e, attr) for attr in pk_attrs): e for e in entities
            }
            return [
                mapping.get(tuple(ident) if not isinstance(ident, tuple) else ident)
                for ident in idents
            ]

    def count(
        self,
        entity: type[_T],
        expressions: Iterable[_ColumnExpressionArgument[bool] | Expression[bool]]
        | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ):
        statement = select(functions.count()).select_from(entity)

        if expressions:
            statement = statement.where(*expressions)
        if filters:
            statement = statement.filter_by(**filters)
        if execution_options:
            statement = statement.execution_options(**execution_options)

        return self.execute(statement).scalar_one()

    def list(
        self,
        entity: type[_T],
        limit: int | None = 100,
        offset: int | None = None,
        order_bys: Iterable[
            _ColumnExpressionOrStrLabelArgument[Any]
            | Column[CriteriaValue]
            | Order[CriteriaValue]
        ]
        | None = None,
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool] | Expression[bool]]
        | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> Sequence[_T]:
        statement = select(entity)

        if limit:
            statement = statement.limit(limit)
        if offset:
            statement = statement.offset(offset)
        if order_bys:
            statement = statement.order_by(*order_bys)
        if options:
            statement = statement.options(*options)
        if expressions:
            statement = statement.where(*expressions)
        if execution_options:
            statement = statement.execution_options(**execution_options)
        if filters:
            statement = statement.filter_by(**filters)

        return self.execute(statement).scalars().all()

    def partitions(
        self,
        entity: type[_T],
        limit: int | None = 100,
        offset: int | None = None,
        size: int | None = 10,
        order_bys: Iterable[
            _ColumnExpressionOrStrLabelArgument[Any]
            | Column[CriteriaValue]
            | Order[CriteriaValue]
        ]
        | None = None,
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool] | Expression[bool]]
        | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> Iterable[Sequence[_T]]:
        statement = select(entity).execution_options(yield_per=size)

        if limit:
            statement = statement.limit(limit)
        if offset:
            statement = statement.offset(offset)
        if order_bys:
            statement = statement.order_by(*order_bys)
        if options:
            statement = statement.options(*options)
        if expressions:
            statement = statement.where(*expressions)
        if execution_options:
            statement = statement.execution_options(**execution_options)
        if filters:
            statement = statement.filter_by(**filters)

        yield from self.execute(statement).scalars().partitions(size)

bulk

bulk(entity: type[_T], idents: Sequence[_PKIdentityArgument], *, options: Sequence[ORMOption] | None = None, execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT) -> list[_T | None]

Bulk version of Session.get. Each element in idents should be exactly the same format as Session.get's ident parameter.

Returns a list of entities in the same order as idents, with None for any ident that was not found.

Source code in arcanus/materia/sqlalchemy/database.py

def bulk(
    self,
    entity: type[_T],
    idents: Sequence[_PKIdentityArgument],
    *,
    options: Sequence[ORMOption] | None = None,
    execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
) -> list[_T | None]:
    """Bulk version of Session.get. Each element in idents should be
    exactly the same format as Session.get's ident parameter.

    Returns a list of entities in the same order as idents, with None
    for any ident that was not found.
    """
    if not idents:
        return []

    mapper = inspect(active_materia.get()[entity])
    pk_columns = mapper.primary_key  # pyright: ignore[reportOptionalMemberAccess]

    if len(pk_columns) == 1:
        # Build the WHERE clause based on single or composite PK
        pk_col = pk_columns[0]
        statement = select(entity).where(pk_col.in_(idents))
    else:
        # Composite PK: use tuple comparison
        # Each ident should be a tuple matching the PK columns
        statement = select(entity).where(tuple_(*pk_columns).in_(idents))

    if options:
        statement = statement.options(*options)
    if execution_options:
        statement = statement.execution_options(**execution_options)

    entities = self.execute(statement).scalars().all()

    # Build mapping from PK value(s) to entity
    if len(pk_columns) == 1:
        pk_attr = pk_columns[0].key
        mapping = {getattr(e, pk_attr): e for e in entities}
        return [mapping.get(ident) for ident in idents]
    else:
        # Composite PK: map tuple of PK values to entity
        pk_attrs = [col.key for col in pk_columns]
        mapping = {
            tuple(getattr(e, attr) for attr in pk_attrs): e for e in entities
        }
        return [
            mapping.get(tuple(ident) if not isinstance(ident, tuple) else ident)
            for ident in idents
        ]

AsyncSession

arcanus.materia.sqlalchemy.database.AsyncSession

Bases: AsyncSession

Source code in arcanus/materia/sqlalchemy/database.py

class AsyncSession(SqlalchemyAsyncSession):
    sync_session_class = Session
    sync_session: Session

    async def __aenter__(self) -> Self:
        self._validation_context_manager = validation_context(self._validation_context)
        self._validation_context_manager.__enter__()
        await super().__aenter__()
        return self

    async def __aexit__(self, type_: Any, value: Any, traceback: Any) -> None:
        if self._validation_context_manager is not None:
            self._validation_context_manager.__exit__(type_, value, traceback)
            self._validation_context_manager = None
        await super().__aexit__(type_, value, traceback)

    @property
    def _validation_context(self) -> ValidationContextT:
        return self.sync_session._validation_context

    @property
    def _validation_context_manager(self) -> ValidateContextGeneratorT | None:
        return self.sync_session._validation_context_manager

    @_validation_context_manager.setter
    def _validation_context_manager(self, value: ValidateContextGeneratorT | None):
        self.sync_session._validation_context_manager = value

    @overload
    async def stream(
        self,
        statement: TypedReturnsRows[tuple[_T]],
        params: _CoreAnyExecuteParams | None = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: _BindArguments | None = None,
        **kw: Any,
    ) -> AsyncResult[tuple[_T]]: ...

    @overload
    async def stream(
        self,
        statement: Executable,
        params: _CoreAnyExecuteParams | None = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: _BindArguments | None = None,
        **kw: Any,
    ) -> AsyncResult[Any]: ...

    async def stream(
        self,
        statement: Executable,
        params: _CoreAnyExecuteParams | None = None,
        *,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        bind_arguments: _BindArguments | None = None,
        **kw: Any,
    ) -> AsyncResult[Any]:
        """Execute a statement and return a streaming
        :class:`AsyncAdaptedResult` object that adapts rows to transmuter types.
        """
        _STREAM_OPTIONS = util.immutabledict({"stream_results": True})

        if execution_options:
            execution_options = util.immutabledict(execution_options).union(
                _STREAM_OPTIONS
            )
        else:
            execution_options = _STREAM_OPTIONS

        result = await util.greenlet_spawn(
            self.sync_session.execute,
            statement,
            params=params,
            execution_options=execution_options,
            bind_arguments=bind_arguments,
            **kw,
        )

        if isinstance(result, AdaptedResult):
            return AsyncAdaptedResult(
                result,
                entities=result.entities,
                force_revalidate=result._force_revalidate,
            )  # pyright: ignore[reportReturnType]

        return AsyncResult(result)

    async def one(
        self,
        entity: type[_T],
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool]] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> _T:
        return await util.greenlet_spawn(
            self.sync_session.one,
            entity,
            options,
            expressions,
            execution_options,
            **filters,
        )

    async def one_or_none(
        self,
        entity: type[_T],
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool]] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> Optional[_T]:
        r = await util.greenlet_spawn(
            self.sync_session.one_or_none,
            entity,
            options,
            expressions,
            execution_options,
            **filters,
        )
        return r

    async def first(
        self,
        entity: type[_T],
        order_bys: Iterable[_ColumnExpressionOrStrLabelArgument[Any]] | None = None,
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool]] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> Optional[_T]:
        return await util.greenlet_spawn(
            self.sync_session.first,
            entity,
            order_bys,
            options,
            expressions,
            execution_options,
            **filters,
        )

    async def bulk(
        self,
        entity: type[_T],
        idents: Sequence[_PKIdentityArgument],
        *,
        options: Sequence[ORMOption] | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
    ) -> list[_T | None]:
        return await util.greenlet_spawn(
            self.sync_session.bulk,
            entity,
            idents,
            options=options,
            execution_options=execution_options,
        )

    async def count(
        self,
        entity: type[_T],
        expressions: Iterable[_ColumnExpressionArgument[bool] | Expression[bool]]
        | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> int:
        return await util.greenlet_spawn(
            self.sync_session.count,
            entity,
            expressions,
            execution_options,
            **filters,
        )

    async def list(
        self,
        entity: type[_T],
        limit: int | None = 100,
        offset: int | None = None,
        order_bys: Iterable[
            _ColumnExpressionOrStrLabelArgument[Any]
            | Column[CriteriaValue]
            | Order[CriteriaValue]
        ]
        | None = None,
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool] | Expression[bool]]
        | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> Sequence[_T]:
        return await util.greenlet_spawn(
            self.sync_session.list,
            entity,
            limit,
            offset,
            order_bys,
            options,
            expressions,
            execution_options,
            **filters,
        )

    async def partitions(
        self,
        entity: type[_T],
        limit: int | None = 100,
        offset: int | None = None,
        size: int | None = 10,
        order_bys: Iterable[
            _ColumnExpressionOrStrLabelArgument[Any]
            | Column[CriteriaValue]
            | Order[CriteriaValue]
        ]
        | None = None,
        options: Iterable[ExecutableOption] | None = None,
        expressions: Iterable[_ColumnExpressionArgument[bool] | Expression[bool]]
        | None = None,
        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
        **filters,
    ) -> AsyncIterator[Sequence[_T]]:
        statement = select(entity)

        if limit:
            statement = statement.limit(limit)
        if offset:
            statement = statement.offset(offset)
        if order_bys:
            statement = statement.order_by(*order_bys)
        if options:
            statement = statement.options(*options)
        if expressions:
            statement = statement.where(*expressions)
        if execution_options:
            statement = statement.execution_options(**execution_options)
        if filters:
            statement = statement.filter_by(**filters)

        async for partition in (
            (
                await self.stream(
                    statement,
                    execution_options=dict(yield_per=size),
                )
            )
            .scalars()
            .partitions(size)
        ):
            yield partition

stream async

stream(statement: TypedReturnsRows[tuple[_T]], params: _CoreAnyExecuteParams | None = None, *, execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT, bind_arguments: _BindArguments | None = None, **kw: Any) -> AsyncResult[tuple[_T]]

stream(statement: Executable, params: _CoreAnyExecuteParams | None = None, *, execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT, bind_arguments: _BindArguments | None = None, **kw: Any) -> AsyncResult[Any]

stream(statement: Executable, params: _CoreAnyExecuteParams | None = None, *, execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT, bind_arguments: _BindArguments | None = None, **kw: Any) -> AsyncResult[Any]

Execute a statement and return a streaming :class:AsyncAdaptedResult object that adapts rows to transmuter types.

Source code in arcanus/materia/sqlalchemy/database.py

async def stream(
    self,
    statement: Executable,
    params: _CoreAnyExecuteParams | None = None,
    *,
    execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
    bind_arguments: _BindArguments | None = None,
    **kw: Any,
) -> AsyncResult[Any]:
    """Execute a statement and return a streaming
    :class:`AsyncAdaptedResult` object that adapts rows to transmuter types.
    """
    _STREAM_OPTIONS = util.immutabledict({"stream_results": True})

    if execution_options:
        execution_options = util.immutabledict(execution_options).union(
            _STREAM_OPTIONS
        )
    else:
        execution_options = _STREAM_OPTIONS

    result = await util.greenlet_spawn(
        self.sync_session.execute,
        statement,
        params=params,
        execution_options=execution_options,
        bind_arguments=bind_arguments,
        **kw,
    )

    if isinstance(result, AdaptedResult):
        return AsyncAdaptedResult(
            result,
            entities=result.entities,
            force_revalidate=result._force_revalidate,
        )  # pyright: ignore[reportReturnType]

    return AsyncResult(result)