# ext/horizontal_shard.py
# Copyright (C) 2005-2024 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
   refering 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 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 Protocol
from ..util.typing import Self

if TYPE_CHECKING:
    from ..engine.base import Connection
    from ..engine.base import Engine
    from ..engine.base import OptionEngine
    from ..engine.result import IteratorResult
    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._typing import _TP
    from ..sql.elements import ClauseElement

__all__ = ["ShardedSession", "ShardedQuery"]

_T = TypeVar("_T", bound=Any)


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"}
            )

        """
        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(  # type: ignore [override]
        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,
) -> Union[Result[_T], IteratorResult[_TP]]:
    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,
    ) -> Union[Result[_T], IteratorResult[_TP]]:
        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:])