shardproxy 1.0.0__tar.gz

shardproxy-1.0.0/PKG-INFO

@@ -0,0 +1,59 @@
+Metadata-Version: 2.3
+Name: shardproxy
+Version: 1.0.0
+Summary: Sharding with SQLAlchemy
+Classifier: License :: OSI Approved :: ISC License (ISCL)
+Requires-Dist: sqlalchemy[asyncio]>=2.0
+Requires-Python: >=3.12
+Description-Content-Type: text/x-rst
+
+shardproxy
+==========
+
+The module consists of two layers:
+
+- ShardProxy class that just manages connections.
+- ops module that implements helper classes for parallelism and result processing.
+
+ShardProxy
+----------
+
+Initialize:
+
+.. code:: python
+
+    SHARD_URLS = [
+        "postgresql+asyncpg://server/shard0",
+        "postgresql+asyncpg://server/shard1",
+    ]
+    proxy = ShardProxy(SHARD_URLS)
+
+Process request:
+
+.. code:: python
+
+    async def fetch_user(user_id: UUID):
+        async with proxy.connect(user_id) as conn:
+            async with conn.begin():
+                stmt = select(User).where(User.id == user_id)
+                res = await conn.execute(stmt)
+                return res.all()
+
+Operations
+----------
+
+Select bunch of IDs:
+
+.. code:: python
+
+    UserRow = Row[tuple[User]]
+
+    class SelectManyUsers(RunOnKeysOperation[UserRow]):
+        async def process(
+            self, conn: AsyncConnection, keys: list[UUID]
+        ) -> ResultRows[UserRow]:
+            stmt = select(User).where(User.id.in_(keys))
+            res = await conn.execute(stmt)
+            return res.all()
+
+    users = await SelectManyUsers().run(proxy)

shardproxy-1.0.0/README.rst

@@ -0,0 +1,50 @@
+shardproxy
+==========
+
+The module consists of two layers:
+
+- ShardProxy class that just manages connections.
+- ops module that implements helper classes for parallelism and result processing.
+
+ShardProxy
+----------
+
+Initialize:
+
+.. code:: python
+
+    SHARD_URLS = [
+        "postgresql+asyncpg://server/shard0",
+        "postgresql+asyncpg://server/shard1",
+    ]
+    proxy = ShardProxy(SHARD_URLS)
+
+Process request:
+
+.. code:: python
+
+    async def fetch_user(user_id: UUID):
+        async with proxy.connect(user_id) as conn:
+            async with conn.begin():
+                stmt = select(User).where(User.id == user_id)
+                res = await conn.execute(stmt)
+                return res.all()
+
+Operations
+----------
+
+Select bunch of IDs:
+
+.. code:: python
+
+    UserRow = Row[tuple[User]]
+
+    class SelectManyUsers(RunOnKeysOperation[UserRow]):
+        async def process(
+            self, conn: AsyncConnection, keys: list[UUID]
+        ) -> ResultRows[UserRow]:
+            stmt = select(User).where(User.id.in_(keys))
+            res = await conn.execute(stmt)
+            return res.all()
+
+    users = await SelectManyUsers().run(proxy)

shardproxy-1.0.0/pyproject.toml

@@ -0,0 +1,53 @@
+[project]
+name = "shardproxy"
+version = "1.0.0"
+description = "Sharding with SQLAlchemy"
+readme = "README.rst"
+classifiers = ["License :: OSI Approved :: ISC License (ISCL)"]
+requires-python = ">=3.12"
+dependencies = ["sqlalchemy[asyncio]>=2.0"]
+
+[dependency-groups]
+dev = [
+    {include-group = "docs"},
+    {include-group = "lint"},
+    {include-group = "test"},
+]
+docs = ["sphinx"]
+lint = ["ruff", "isort", "mypy"]
+test = [
+    "asyncpg",
+    "psycopg[binary,pool]",
+    "pytest",
+    "pytest-asyncio",
+    "pytest-cov",
+    "coverage",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.1,<0.12.0"]
+build-backend = "uv_build"
+
+[tool.mypy]
+python_version = "3.12"
+strict = true
+disallow_any_decorated = true
+disallow_any_unimported = true
+#disallow_any_explicit = true
+#disallow_any_expr = true
+
+[tool.isort]
+known_first_party = ["shardproxy"]
+known_local_folder = ["helpers"]
+multi_line_output = 5
+balanced_wrapping = true
+include_trailing_comma = true
+atomic = true
+quiet = true
+
+[[tool.uv.index]]
+name = "testpypi"
+url = "https://test.pypi.org/simple/"
+publish-url = "https://test.pypi.org/legacy/"
+explicit = true
+

shardproxy-1.0.0/src/shardproxy/__init__.py

@@ -0,0 +1,24 @@
+"""Sharding with SQLAlchemy."""
+
+from .ops import (
+    RunOnAllOperation,
+    RunOnAnyOperation,
+    RunOnArgsOperation,
+    RunOnKeysOperation,
+    RunOnOneOperation,
+)
+from .proxy import ShardProxy
+from .types import AsyncConnection, AsyncEngine, DBTask, ResultRows
+
+__all__ = (
+    "ShardProxy",
+    "AsyncEngine",
+    "AsyncConnection",
+    "ResultRows",
+    "DBTask",
+    "RunOnAllOperation",
+    "RunOnKeysOperation",
+    "RunOnArgsOperation",
+    "RunOnAnyOperation",
+    "RunOnOneOperation",
+)

shardproxy-1.0.0/src/shardproxy/ops.py

@@ -0,0 +1,234 @@
+"""Helper classes for multi-shard operations."""
+
+import asyncio
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from uuid import UUID
+
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+
+from .proxy import ShardProxy
+from .types import DBTask, ResultRows
+
+__all__ = (
+    "RunOnKeysOperation",
+    "RunOnArgsOperation",
+    "RunOnAllOperation",
+    "RunOnOneOperation",
+    "RunOnAnyOperation",
+)
+
+
+class RunOnKeysOperation[TRow](ABC):
+    """Runs operation over given keys.
+
+    Usage::
+
+        class MyOp(RunOnKeysOperation[MyRow]):
+            async def process(
+                self, conn: AsyncConnection, keys: list[UUID]
+            ) -> ResultRows[MyRow]:
+                return ...
+
+    """
+
+    #: Whether connection is in transaction context when given to .process().
+    USE_TRANSACTION = True
+
+    @abstractmethod
+    async def process(
+        self, conn: AsyncConnection, keys: list[UUID]
+    ) -> ResultRows[TRow]:
+        """Business logic to execute on each shard."""
+        raise NotImplementedError
+
+    async def run(
+        self, proxy: ShardProxy, shard_keys: Sequence[UUID]
+    ) -> ResultRows[TRow]:
+        """Split keys for each shard and run .process() on them."""
+
+        tasks: list[DBTask[ResultRows[TRow]]] = []
+        for engine, keys in proxy.spread_keys(shard_keys):
+            task = asyncio.create_task(self._handle_process(engine, keys))
+            tasks.append(task)
+
+        return await self._collect(tasks)
+
+    async def _handle_process(
+        self, engine: AsyncEngine, keys: list[UUID]
+    ) -> ResultRows[TRow]:
+        async with engine.connect() as conn:
+            if self.USE_TRANSACTION:
+                async with conn.begin():
+                    return await self.process(conn, keys)
+            else:
+                return await self.process(conn, keys)
+
+    async def _collect(self, tasks: list[DBTask[ResultRows[TRow]]]) -> ResultRows[TRow]:
+        shard_rows = await asyncio.gather(*tasks)
+        res: list[TRow] = []
+        for rows in shard_rows:
+            res.extend(rows)
+        return res
+
+
+class RunOnArgsOperation[TRow, TArg](ABC):
+    """Runs operation over given arguments.
+
+    Usage::
+
+        class MyOp(RunOnArgsOperation[MyRow, MyArg]):
+            async def process(
+                self, conn: AsyncConnection, args: dict[UUID, MyArg]
+            ) -> ResultRows[MyRow]:
+                return ...
+
+    """
+
+    #: Whether connection is in transaction context when given to .process().
+    USE_TRANSACTION = True
+
+    @abstractmethod
+    async def process(
+        self, conn: AsyncConnection, args: dict[UUID, TArg]
+    ) -> ResultRows[TRow]:
+        """Business logic to execute on each shard."""
+        raise NotImplementedError
+
+    async def run(self, proxy: ShardProxy, args: dict[UUID, TArg]) -> ResultRows[TRow]:
+        """Split args for each shard and run .process() with them."""
+
+        tasks: list[DBTask[ResultRows[TRow]]] = []
+        for engine, shard_args in proxy.spread_args(args):
+            task = asyncio.create_task(self._handle_process(engine, shard_args))
+            tasks.append(task)
+
+        return await self._collect(tasks)
+
+    async def _handle_process(
+        self, engine: AsyncEngine, args: dict[UUID, TArg]
+    ) -> ResultRows[TRow]:
+        async with engine.connect() as conn:
+            if self.USE_TRANSACTION:
+                async with conn.begin():
+                    return await self.process(conn, args)
+            else:
+                return await self.process(conn, args)
+
+    async def _collect(self, tasks: list[DBTask[ResultRows[TRow]]]) -> ResultRows[TRow]:
+        shard_rows = await asyncio.gather(*tasks)
+        res: list[TRow] = []
+        for rows in shard_rows:
+            res.extend(rows)
+        return res
+
+
+class RunOnAllOperation[TRow](ABC):
+    """Runs operation over all shards.
+
+    Usage::
+
+        class MyOp(RunOnAllOperation[MyRow]):
+            async def process(
+                self, conn: AsyncConnection
+            ) -> ResultRows[MyRow]:
+                return ...
+
+    """
+
+    #: Whether connection is in transaction context when given to .process().
+    USE_TRANSACTION = True
+
+    @abstractmethod
+    async def process(self, conn: AsyncConnection) -> ResultRows[TRow]:
+        """Business logic to execute on each shard."""
+        raise NotImplementedError
+
+    async def run(self, proxy: ShardProxy) -> ResultRows[TRow]:
+        """Loop over all shards, execute .process() on each."""
+
+        tasks: list[DBTask[ResultRows[TRow]]] = []
+        for engine in proxy.shards:
+            task = asyncio.create_task(self._handle_process(engine))
+            tasks.append(task)
+
+        return await self._collect(tasks)
+
+    async def _handle_process(self, engine: AsyncEngine) -> ResultRows[TRow]:
+        async with engine.connect() as conn:
+            if self.USE_TRANSACTION:
+                async with conn.begin():
+                    return await self.process(conn)
+            else:
+                return await self.process(conn)
+
+    async def _collect(self, tasks: list[DBTask[ResultRows[TRow]]]) -> ResultRows[TRow]:
+        shard_rows = await asyncio.gather(*tasks)
+        res: list[TRow] = []
+        for rows in shard_rows:
+            res.extend(rows)
+        return res
+
+
+class RunOnOneOperation[TRow](ABC):
+    """Runs operation in single shard.
+
+    Usage::
+
+        class MyOp(RunOnOneOperation[MyRow]):
+            async def process(
+                self, conn: AsyncConnection, shard_key: UUID
+            ) -> ResultRows[MyRow]:
+                return ...
+
+    """
+
+    #: Whether connection is in transaction context when given to .process().
+    USE_TRANSACTION = True
+
+    @abstractmethod
+    async def process(self, conn: AsyncConnection, shard_key: UUID) -> ResultRows[TRow]:
+        """Business logic to execute on shard."""
+        raise NotImplementedError
+
+    async def run(self, proxy: ShardProxy, shard_key: UUID) -> ResultRows[TRow]:
+        """Pick shard and run .process() on it."""
+
+        async with proxy.connect(shard_key) as conn:
+            if self.USE_TRANSACTION:
+                async with conn.begin():
+                    return await self.process(conn, shard_key)
+            else:
+                return await self.process(conn, shard_key)
+
+
+class RunOnAnyOperation[TRow](ABC):
+    """Runs operation in single random shard.
+
+    Usage::
+
+        class MyOp(RunOnAnyOperation[MyRow]):
+            async def process(
+                self, conn: AsyncConnection
+            ) -> ResultRows[MyRow]:
+                return ...
+
+    """
+
+    #: Whether connection is in transaction context when given to .process().
+    USE_TRANSACTION = True
+
+    @abstractmethod
+    async def process(self, conn: AsyncConnection) -> ResultRows[TRow]:
+        """Business logic to execute on shard."""
+        raise NotImplementedError
+
+    async def run(self, proxy: ShardProxy) -> ResultRows[TRow]:
+        """Pick random shard and run .process() on it."""
+
+        async with proxy.connect_any() as conn:
+            if self.USE_TRANSACTION:
+                async with conn.begin():
+                    return await self.process(conn)
+            else:
+                return await self.process(conn)

shardproxy-1.0.0/src/shardproxy/proxy.py

@@ -0,0 +1,136 @@
+"""Connection manager for sharding."""
+
+from collections.abc import Mapping, Sequence
+from random import choice
+from typing import Unpack
+from uuid import UUID
+
+from sqlalchemy import URL
+from sqlalchemy.ext.asyncio import (
+    AsyncConnection,
+    AsyncEngine,
+    create_async_engine,
+)
+
+from .types import AsyncEngineParams
+
+__all__ = ("ShardProxy",)
+
+
+class ShardProxy:
+    """Manages AsyncEngine mapping for shards.
+
+    Takes list of shard URLs and generic create_async_engine()/create_engine()
+    keyword arguments.
+    """
+
+    _shards: list[AsyncEngine]
+
+    def __init__(
+        self, shard_urls: Sequence[str | URL], **kwargs: Unpack[AsyncEngineParams]
+    ) -> None:
+        self._shards = []
+        for url in shard_urls:
+            engine = create_async_engine(url, **kwargs)
+            self._shards.append(engine)
+
+    def _get_shard(self, shard_key: UUID) -> int:
+        return shard_key.node % len(self._shards)
+
+    @property
+    def shards(self) -> Sequence[AsyncEngine]:
+        """Read-only list of engines per shard.
+
+        Usage::
+
+            # run-on-all: serial mode
+            for engine in proxy.shards:
+                async with engine.connect() as conn:
+                    async with conn.begin():
+                        await conn.execute()
+
+            # run-on-all: parallel mode
+            async with asyncio.TaskGroup() as tg:
+                for engine in proxy.shards:
+                    tg.create_task(process(engine, items))
+        """
+        return self._shards
+
+    def connect(self, shard_key: UUID) -> AsyncConnection:
+        """Returns connection for specific shard as context manager.
+
+        Usage::
+
+            async with proxy.connect(key) as conn:
+                async with conn.begin():
+                    ...
+        """
+        idx = self._get_shard(shard_key)
+        return self._shards[idx].connect()
+
+    def connect_any[T](self) -> AsyncConnection:
+        """Returns connection to random shard as context manager.
+
+        Usage::
+
+            async with proxy.connect_any() as conn:
+                async with conn.begin():
+                    ...
+        """
+        engine = choice(self._shards)
+        return engine.connect()
+
+    def spread_args[A](
+        self, args: Mapping[UUID, A]
+    ) -> Sequence[tuple[AsyncEngine, dict[UUID, A]]]:
+        """Spread keys with arguments over shards.
+
+        Usage::
+
+            # serial mode
+            for engine, shard_items in proxy.split_args(all_items):
+                async with engine.connect() as conn:
+                    async with conn.begin():
+                        await conn.execute()
+
+            # parallel mode
+            async with asyncio.TaskGroup() as tg:
+                for engine, items in proxy.split_args(all_items):
+                    tg.create_task(process(engine, items))
+        """
+        split: dict[int, dict[UUID, A]] = {}
+        for key, arg in args.items():
+            nr = self._get_shard(key)
+            target = split.get(nr)
+            if target is None:
+                split[nr] = {key: arg}
+            else:
+                target[key] = arg
+        return [(self._shards[nr], target) for nr, target in split.items()]
+
+    def spread_keys(
+        self, keys: Sequence[UUID]
+    ) -> Sequence[tuple[AsyncEngine, list[UUID]]]:
+        """Spread keys over shards.
+
+        Usage::
+
+            # serial mode
+            for engine, shard_keys in proxy.split_keys(all_keys):
+                async with engine.connect() as conn:
+                    async with conn.begin():
+                        await conn.execute()
+
+            # parallel mode
+            async with asyncio.TaskGroup() as tg:
+                for engine, shard_keys in proxy.split_keys(all_keys):
+                    tg.create_task(process(engine, shard_keys))
+        """
+        split: dict[int, set[UUID]] = {}
+        for key in keys:
+            nr = self._get_shard(key)
+            target = split.get(nr)
+            if target is None:
+                target = split[nr] = set()
+            target.add(key)
+        return [(self._shards[nr], list(target)) for nr, target in split.items()]

shardproxy-1.0.0/src/shardproxy/types.py

@@ -0,0 +1,68 @@
+"""Common types"""
+
+import asyncio
+from collections.abc import Sequence
+from typing import Any, Callable, Literal, Type, TypedDict
+
+from sqlalchemy.engine.interfaces import (
+    CoreExecuteOptionsParameter,
+    IsolationLevel,
+)
+from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine
+from sqlalchemy.pool import Pool
+
+__all__ = (
+    "AsyncConnection",
+    "AsyncEngine",
+    "AsyncEngineParams",
+    "ResultRows",
+    "DBTask",
+)
+
+
+type ResultRows[TRow] = Sequence[TRow]
+
+type DBTask[TResult] = asyncio.Task[TResult]
+
+
+class AsyncEngineParams(TypedDict, total=False):
+    """Helper struct to describe common parameters for create_async_engine.
+
+    It may be incomplete - arguments from here are
+    passed to AsyncEngine, Dialect and Pool.
+
+    Some values are deliberately hidden as they will not work with ShardProxy.
+    """
+
+    # async_creator - does not work with sharding
+    # connect_args - does not work with sharding
+    echo: bool | Literal["debug"] | None
+    echo_pool: bool | Literal["debug"] | None
+    enable_from_linting: bool
+    execution_options: CoreExecuteOptionsParameter
+    hide_parameters: bool
+    insertmanyvalues_page_size: int
+    isolation_level: IsolationLevel
+    json_deserializer: Callable[..., Any]
+    json_serializer: Callable[..., Any]
+    label_length: int | None
+    logging_name: str
+    max_identifier_length: int
+    max_overflow: int
+    # module - does not work with sharding
+    paramstyle: Literal[
+        "qmark", "numeric", "named", "format", "pyformat", "numeric_dollar"
+    ]
+    # pool - does not work with sharding
+    poolclass: Type[Pool]
+    pool_logging_name: str
+    pool_pre_ping: bool
+    pool_size: int
+    pool_recycle: int
+    pool_reset_on_return: Literal["rollback", "commit"] | None
+    pool_timeout: int
+    pool_use_iifo: bool
+    plugins: list[str]
+    query_cache_size: int
+    skip_autocommit_rollback: bool
+    use_insertmanyvalues: bool