SQLAlchemy 2.1.0b2__cp313-cp313t-win_arm64.whl

sqlalchemy/ext/horizontal_shard.py

@@ -0,0 +1,481 @@
+# ext/horizontal_shard.py
+# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of SQLAlchemy and is released under
+# the MIT License: https://www.opensource.org/licenses/mit-license.php
+
+"""Horizontal sharding support.
+
+Defines a rudimental 'horizontal sharding' system which allows a Session to
+distribute queries and persistence operations across multiple databases.
+
+For a usage example, see the :ref:`examples_sharding` example included in
+the source distribution.
+
+.. deepalchemy:: The horizontal sharding extension is an advanced feature,
+   involving a complex statement -> database interaction as well as
+   use of semi-public APIs for non-trivial cases.   Simpler approaches to
+   referring to multiple database "shards", most commonly using a distinct
+   :class:`_orm.Session` per "shard", should always be considered first
+   before using this more complex and less-production-tested system.
+
+
+
+"""
+from __future__ import annotations
+
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import Iterable
+from typing import Optional
+from typing import Protocol
+from typing import Tuple
+from typing import Type
+from typing import TYPE_CHECKING
+from typing import TypeVar
+from typing import Union
+
+from .. import event
+from .. import exc
+from .. import inspect
+from .. import util
+from ..orm import PassiveFlag
+from ..orm._typing import OrmExecuteOptionsParameter
+from ..orm.interfaces import ORMOption
+from ..orm.mapper import Mapper
+from ..orm.query import Query
+from ..orm.session import _BindArguments
+from ..orm.session import _PKIdentityArgument
+from ..orm.session import Session
+from ..util.typing import Self
+from ..util.typing import TupleAny
+from ..util.typing import TypeVarTuple
+from ..util.typing import Unpack
+
+
+if TYPE_CHECKING:
+    from ..engine.base import Connection
+    from ..engine.base import Engine
+    from ..engine.base import OptionEngine
+    from ..engine.result import Result
+    from ..orm import LoaderCallableStatus
+    from ..orm._typing import _O
+    from ..orm.bulk_persistence import _BulkUDCompileState
+    from ..orm.context import QueryContext
+    from ..orm.session import _EntityBindKey
+    from ..orm.session import _SessionBind
+    from ..orm.session import ORMExecuteState
+    from ..orm.state import InstanceState
+    from ..sql import Executable
+    from ..sql.elements import ClauseElement
+
+__all__ = ["ShardedSession", "ShardedQuery"]
+
+_T = TypeVar("_T", bound=Any)
+_Ts = TypeVarTuple("_Ts")
+
+
+ShardIdentifier = str
+
+
+class ShardChooser(Protocol):
+    def __call__(
+        self,
+        mapper: Optional[Mapper[_T]],
+        instance: Any,
+        clause: Optional[ClauseElement],
+    ) -> Any: ...
+
+
+class IdentityChooser(Protocol):
+    def __call__(
+        self,
+        mapper: Mapper[_T],
+        primary_key: _PKIdentityArgument,
+        *,
+        lazy_loaded_from: Optional[InstanceState[Any]],
+        execution_options: OrmExecuteOptionsParameter,
+        bind_arguments: _BindArguments,
+        **kw: Any,
+    ) -> Any: ...
+
+
+class ShardedQuery(Query[_T]):
+    """Query class used with :class:`.ShardedSession`.
+
+    .. legacy:: The :class:`.ShardedQuery` is a subclass of the legacy
+       :class:`.Query` class.   The :class:`.ShardedSession` now supports
+       2.0 style execution via the :meth:`.ShardedSession.execute` method.
+
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        assert isinstance(self.session, ShardedSession)
+
+        self.identity_chooser = self.session.identity_chooser
+        self.execute_chooser = self.session.execute_chooser
+        self._shard_id = None
+
+    def set_shard(self, shard_id: ShardIdentifier) -> Self:
+        """Return a new query, limited to a single shard ID.
+
+        All subsequent operations with the returned query will
+        be against the single shard regardless of other state.
+
+        The shard_id can be passed for a 2.0 style execution to the
+        bind_arguments dictionary of :meth:`.Session.execute`::
+
+            results = session.execute(stmt, bind_arguments={"shard_id": "my_shard"})
+
+        """  # noqa: E501
+        return self.execution_options(_sa_shard_id=shard_id)
+
+
+class ShardedSession(Session):
+    shard_chooser: ShardChooser
+    identity_chooser: IdentityChooser
+    execute_chooser: Callable[[ORMExecuteState], Iterable[Any]]
+
+    def __init__(
+        self,
+        shard_chooser: ShardChooser,
+        identity_chooser: Optional[IdentityChooser] = None,
+        execute_chooser: Optional[
+            Callable[[ORMExecuteState], Iterable[Any]]
+        ] = None,
+        shards: Optional[Dict[str, Any]] = None,
+        query_cls: Type[Query[_T]] = ShardedQuery,
+        *,
+        id_chooser: Optional[
+            Callable[[Query[_T], Iterable[_T]], Iterable[Any]]
+        ] = None,
+        query_chooser: Optional[Callable[[Executable], Iterable[Any]]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """Construct a ShardedSession.
+
+        :param shard_chooser: A callable which, passed a Mapper, a mapped
+          instance, and possibly a SQL clause, returns a shard ID.  This id
+          may be based off of the attributes present within the object, or on
+          some round-robin scheme. If the scheme is based on a selection, it
+          should set whatever state on the instance to mark it in the future as
+          participating in that shard.
+
+        :param identity_chooser: A callable, passed a Mapper and primary key
+         argument, which should return a list of shard ids where this
+         primary key might reside.
+
+          .. versionchanged:: 2.0  The ``identity_chooser`` parameter
+             supersedes the ``id_chooser`` parameter.
+
+        :param execute_chooser: For a given :class:`.ORMExecuteState`,
+          returns the list of shard_ids
+          where the query should be issued.  Results from all shards returned
+          will be combined together into a single listing.
+
+          .. versionchanged:: 1.4  The ``execute_chooser`` parameter
+             supersedes the ``query_chooser`` parameter.
+
+        :param shards: A dictionary of string shard names
+          to :class:`~sqlalchemy.engine.Engine` objects.
+
+        """
+        super().__init__(query_cls=query_cls, **kwargs)
+
+        event.listen(
+            self, "do_orm_execute", execute_and_instances, retval=True
+        )
+        self.shard_chooser = shard_chooser
+
+        if id_chooser:
+            _id_chooser = id_chooser
+            util.warn_deprecated(
+                "The ``id_chooser`` parameter is deprecated; "
+                "please use ``identity_chooser``.",
+                "2.0",
+            )
+
+            def _legacy_identity_chooser(
+                mapper: Mapper[_T],
+                primary_key: _PKIdentityArgument,
+                *,
+                lazy_loaded_from: Optional[InstanceState[Any]],
+                execution_options: OrmExecuteOptionsParameter,
+                bind_arguments: _BindArguments,
+                **kw: Any,
+            ) -> Any:
+                q = self.query(mapper)
+                if lazy_loaded_from:
+                    q = q._set_lazyload_from(lazy_loaded_from)
+                return _id_chooser(q, primary_key)
+
+            self.identity_chooser = _legacy_identity_chooser
+        elif identity_chooser:
+            self.identity_chooser = identity_chooser
+        else:
+            raise exc.ArgumentError(
+                "identity_chooser or id_chooser is required"
+            )
+
+        if query_chooser:
+            _query_chooser = query_chooser
+            util.warn_deprecated(
+                "The ``query_chooser`` parameter is deprecated; "
+                "please use ``execute_chooser``.",
+                "1.4",
+            )
+            if execute_chooser:
+                raise exc.ArgumentError(
+                    "Can't pass query_chooser and execute_chooser "
+                    "at the same time."
+                )
+
+            def _default_execute_chooser(
+                orm_context: ORMExecuteState,
+            ) -> Iterable[Any]:
+                return _query_chooser(orm_context.statement)
+
+            if execute_chooser is None:
+                execute_chooser = _default_execute_chooser
+
+        if execute_chooser is None:
+            raise exc.ArgumentError(
+                "execute_chooser or query_chooser is required"
+            )
+        self.execute_chooser = execute_chooser
+        self.__shards: Dict[ShardIdentifier, _SessionBind] = {}
+        if shards is not None:
+            for k in shards:
+                self.bind_shard(k, shards[k])
+
+    def _identity_lookup(
+        self,
+        mapper: Mapper[_O],
+        primary_key_identity: Union[Any, Tuple[Any, ...]],
+        identity_token: Optional[Any] = None,
+        passive: PassiveFlag = PassiveFlag.PASSIVE_OFF,
+        lazy_loaded_from: Optional[InstanceState[Any]] = None,
+        execution_options: OrmExecuteOptionsParameter = util.EMPTY_DICT,
+        bind_arguments: Optional[_BindArguments] = None,
+        **kw: Any,
+    ) -> Union[Optional[_O], LoaderCallableStatus]:
+        """override the default :meth:`.Session._identity_lookup` method so
+        that we search for a given non-token primary key identity across all
+        possible identity tokens (e.g. shard ids).
+
+        .. versionchanged:: 1.4  Moved :meth:`.Session._identity_lookup` from
+           the :class:`_query.Query` object to the :class:`.Session`.
+
+        """
+
+        if identity_token is not None:
+            obj = super()._identity_lookup(
+                mapper,
+                primary_key_identity,
+                identity_token=identity_token,
+                **kw,
+            )
+
+            return obj
+        else:
+            for shard_id in self.identity_chooser(
+                mapper,
+                primary_key_identity,
+                lazy_loaded_from=lazy_loaded_from,
+                execution_options=execution_options,
+                bind_arguments=dict(bind_arguments) if bind_arguments else {},
+            ):
+                obj2 = super()._identity_lookup(
+                    mapper,
+                    primary_key_identity,
+                    identity_token=shard_id,
+                    lazy_loaded_from=lazy_loaded_from,
+                    **kw,
+                )
+                if obj2 is not None:
+                    return obj2
+
+            return None
+
+    def _choose_shard_and_assign(
+        self,
+        mapper: Optional[_EntityBindKey[_O]],
+        instance: Any,
+        **kw: Any,
+    ) -> Any:
+        if instance is not None:
+            state = inspect(instance)
+            if state.key:
+                token = state.key[2]
+                assert token is not None
+                return token
+            elif state.identity_token:
+                return state.identity_token
+
+        assert isinstance(mapper, Mapper)
+        shard_id = self.shard_chooser(mapper, instance, **kw)
+        if instance is not None:
+            state.identity_token = shard_id
+        return shard_id
+
+    def connection_callable(
+        self,
+        mapper: Optional[Mapper[_T]] = None,
+        instance: Optional[Any] = None,
+        shard_id: Optional[ShardIdentifier] = None,
+        **kw: Any,
+    ) -> Connection:
+        """Provide a :class:`_engine.Connection` to use in the unit of work
+        flush process.
+
+        """
+
+        if shard_id is None:
+            shard_id = self._choose_shard_and_assign(mapper, instance)
+
+        if self.in_transaction():
+            trans = self.get_transaction()
+            assert trans is not None
+            return trans.connection(mapper, shard_id=shard_id)
+        else:
+            bind = self.get_bind(
+                mapper=mapper, shard_id=shard_id, instance=instance
+            )
+
+            if isinstance(bind, Engine):
+                return bind.connect(**kw)
+            else:
+                assert isinstance(bind, Connection)
+                return bind
+
+    def get_bind(
+        self,
+        mapper: Optional[_EntityBindKey[_O]] = None,
+        *,
+        shard_id: Optional[ShardIdentifier] = None,
+        instance: Optional[Any] = None,
+        clause: Optional[ClauseElement] = None,
+        **kw: Any,
+    ) -> _SessionBind:
+        if shard_id is None:
+            shard_id = self._choose_shard_and_assign(
+                mapper, instance=instance, clause=clause
+            )
+            assert shard_id is not None
+        return self.__shards[shard_id]
+
+    def bind_shard(
+        self, shard_id: ShardIdentifier, bind: Union[Engine, OptionEngine]
+    ) -> None:
+        self.__shards[shard_id] = bind
+
+
+class set_shard_id(ORMOption):
+    """a loader option for statements to apply a specific shard id to the
+    primary query as well as for additional relationship and column
+    loaders.
+
+    The :class:`_horizontal.set_shard_id` option may be applied using
+    the :meth:`_sql.Executable.options` method of any executable statement::
+
+        stmt = (
+            select(MyObject)
+            .where(MyObject.name == "some name")
+            .options(set_shard_id("shard1"))
+        )
+
+    Above, the statement when invoked will limit to the "shard1" shard
+    identifier for the primary query as well as for all relationship and
+    column loading strategies, including eager loaders such as
+    :func:`_orm.selectinload`, deferred column loaders like :func:`_orm.defer`,
+    and the lazy relationship loader :func:`_orm.lazyload`.
+
+    In this way, the :class:`_horizontal.set_shard_id` option has much wider
+    scope than using the "shard_id" argument within the
+    :paramref:`_orm.Session.execute.bind_arguments` dictionary.
+
+
+    .. versionadded:: 2.0.0
+
+    """
+
+    __slots__ = ("shard_id", "propagate_to_loaders")
+
+    def __init__(
+        self, shard_id: ShardIdentifier, propagate_to_loaders: bool = True
+    ):
+        """Construct a :class:`_horizontal.set_shard_id` option.
+
+        :param shard_id: shard identifier
+        :param propagate_to_loaders: if left at its default of ``True``, the
+         shard option will take place for lazy loaders such as
+         :func:`_orm.lazyload` and :func:`_orm.defer`; if False, the option
+         will not be propagated to loaded objects. Note that :func:`_orm.defer`
+         always limits to the shard_id of the parent row in any case, so the
+         parameter only has a net effect on the behavior of the
+         :func:`_orm.lazyload` strategy.
+
+        """
+        self.shard_id = shard_id
+        self.propagate_to_loaders = propagate_to_loaders
+
+
+def execute_and_instances(
+    orm_context: ORMExecuteState,
+) -> Result[Unpack[TupleAny]]:
+    active_options: Union[
+        None,
+        QueryContext.default_load_options,
+        Type[QueryContext.default_load_options],
+        _BulkUDCompileState.default_update_options,
+        Type[_BulkUDCompileState.default_update_options],
+    ]
+
+    if orm_context.is_select:
+        active_options = orm_context.load_options
+
+    elif orm_context.is_update or orm_context.is_delete:
+        active_options = orm_context.update_delete_options
+    else:
+        active_options = None
+
+    session = orm_context.session
+    assert isinstance(session, ShardedSession)
+
+    def iter_for_shard(
+        shard_id: ShardIdentifier,
+    ) -> Result[Unpack[TupleAny]]:
+        bind_arguments = dict(orm_context.bind_arguments)
+        bind_arguments["shard_id"] = shard_id
+
+        orm_context.update_execution_options(identity_token=shard_id)
+        return orm_context.invoke_statement(bind_arguments=bind_arguments)
+
+    for orm_opt in orm_context._non_compile_orm_options:
+        # TODO: if we had an ORMOption that gets applied at ORM statement
+        # execution time, that would allow this to be more generalized.
+        # for now just iterate and look for our options
+        if isinstance(orm_opt, set_shard_id):
+            shard_id = orm_opt.shard_id
+            break
+    else:
+        if active_options and active_options._identity_token is not None:
+            shard_id = active_options._identity_token
+        elif "_sa_shard_id" in orm_context.execution_options:
+            shard_id = orm_context.execution_options["_sa_shard_id"]
+        elif "shard_id" in orm_context.bind_arguments:
+            shard_id = orm_context.bind_arguments["shard_id"]
+        else:
+            shard_id = None
+
+    if shard_id is not None:
+        return iter_for_shard(shard_id)
+    else:
+        partial = []
+        for shard_id in session.execute_chooser(orm_context):
+            result_ = iter_for_shard(shard_id)
+            partial.append(result_)
+        return partial[0].merge(*partial[1:])