from typing import Any
from typing import Optional

from ..orm.query import Query
from ..orm.session import Session

class ShardedQuery(Query):
    id_chooser: Any = ...
    query_chooser: Any = ...
    execute_chooser: Any = ...
    def __init__(self, *args: Any, **kwargs: Any) -> None: ...
    def set_shard(self, shard_id: Any): ...

class ShardedSession(Session):
    shard_chooser: Any = ...
    id_chooser: Any = ...
    execute_chooser: Any = ...
    query_chooser: Any = ...
    def __init__(
        self,
        shard_chooser: Any,
        id_chooser: Any,
        execute_chooser: Optional[Any] = ...,
        shards: Optional[Any] = ...,
        query_cls: Any = ...,
        **kwargs: Any,
    ): ...
    def connection_callable(
        self,
        mapper: Optional[Any] = ...,
        instance: Optional[Any] = ...,
        shard_id: Optional[Any] = ...,
        **kwargs: Any,
    ): ...
    def get_bind(
        self,
        mapper: Optional[Any] = ...,
        shard_id: Optional[Any] = ...,
        instance: Optional[Any] = ...,
        clause: Optional[Any] = ...,
        **kw: Any,
    ): ...
    def bind_shard(self, shard_id: Any, bind: Any) -> None: ...
