feat(service): add service composition primitives

Author: cofin

advanced_alchemy/service/__init__.py

@@ -19,7 +19,9 @@
     SQLAlchemySyncRepositoryReadService,
     SQLAlchemySyncRepositoryService,
 )
+from advanced_alchemy.service._typing import ServiceProvider, ServiceWithSession
 from advanced_alchemy.service._util import ResultConverter, find_filter
+from advanced_alchemy.service.composition import ServiceComposition
 from advanced_alchemy.service.pagination import OffsetPagination
 from advanced_alchemy.typing import ATTRS_INSTALLED
 from advanced_alchemy.utils.serialization import (
@@ -77,6 +79,9 @@
     "SQLAlchemySyncQueryService",
     "SQLAlchemySyncRepositoryReadService",
     "SQLAlchemySyncRepositoryService",
+    "ServiceComposition",
+    "ServiceProvider",
+    "ServiceWithSession",
     "SupportedSchemaModel",
     "fields",
     "find_filter",

advanced_alchemy/service/_typing.py

@@ -0,0 +1,23 @@
+"""Service-side typing surface."""
+
+from collections.abc import AsyncGenerator, Callable
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
+
+from typing_extensions import TypeAlias, TypeVar
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+
+__all__ = ("ServiceProvider", "ServiceWithSession")
+
+T = TypeVar("T")
+
+ServiceProvider: TypeAlias = Callable[["AsyncSession"], AsyncGenerator[T, None]]
+"""Callable that yields a service instance for a session."""
+
+
+@runtime_checkable
+class ServiceWithSession(Protocol):
+    """Structural protocol for services constructible with ``session=``."""
+
+    def __init__(self, *, session: "AsyncSession") -> None: ...

advanced_alchemy/service/composition.py

@@ -0,0 +1,80 @@
+"""Framework-agnostic service composition primitives.
+
+Example:
+    Compose providers that yield services for a caller-owned async session:
+
+    .. code-block:: python
+
+        async with (
+            ServiceComposition.starting_with(provide_users_service)
+            .add(provide_roles_service)
+            .open(session=db_session)
+        ) as (users, roles):
+            await users.create(data)
+            await roles.assign_default_role(data)
+"""
+
+from contextlib import AsyncExitStack, asynccontextmanager
+from typing import TYPE_CHECKING, Any, Generic, cast
+
+from typing_extensions import TypeVar, TypeVarTuple, Unpack
+
+from advanced_alchemy.service._typing import ServiceProvider
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    from sqlalchemy.ext.asyncio import AsyncSession
+
+__all__ = ("ServiceComposition", "ServiceProvider")
+
+T = TypeVar("T")
+Ts = TypeVarTuple("Ts")
+
+
+@asynccontextmanager
+async def _aclose_generator(
+    generator: "AsyncGenerator[Any, None]",
+) -> "AsyncGenerator[AsyncGenerator[Any, None], None]":
+    # Python 3.9 lacks contextlib.aclosing; replace this helper when 3.9 support is dropped.
+    try:
+        yield generator
+    finally:
+        await generator.aclose()
+
+
+class ServiceComposition(Generic[Unpack[Ts]]):
+    """Builder for composing services with a shared async session.
+
+    Use :meth:`starting_with` to seed the composition, then chain
+    :meth:`add` calls to preserve each service's positional type.
+    """
+
+    __slots__ = ("_providers",)
+
+    def __init__(self, providers: "tuple[ServiceProvider[Any], ...]" = ()) -> None:
+        self._providers = providers
+
+    @classmethod
+    def starting_with(cls, provider: "ServiceProvider[T]") -> "ServiceComposition[T]":
+        """Begin a composition with the first provider."""
+        return cast("ServiceComposition[T]", cls((provider,)))
+
+    def add(self, provider: "ServiceProvider[T]") -> "ServiceComposition[Unpack[Ts], T]":
+        """Append a provider, extending the service tuple by one position."""
+        return cast("ServiceComposition[Unpack[Ts], T]", ServiceComposition((*self._providers, provider)))
+
+    @asynccontextmanager
+    async def open(self, *, session: "AsyncSession") -> "AsyncGenerator[tuple[Unpack[Ts]], None]":
+        """Enter all providers on ``session`` and yield their services.
+
+        The session lifecycle remains owned by the caller. Entered provider
+        generators are closed in reverse order on normal exit, user exception,
+        or partial-entry failure.
+        """
+        services: list[Any] = []
+        async with AsyncExitStack() as stack:
+            for provider in self._providers:
+                generator = await stack.enter_async_context(_aclose_generator(provider(session)))
+                services.append(await generator.__anext__())
+            yield cast("tuple[Unpack[Ts]]", tuple(services))

tests/unit/test_service/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for service helpers."""

tests/unit/test_service/test_composition.py

@@ -0,0 +1,130 @@
+"""Tests for service composition primitives."""
+
+from collections.abc import AsyncGenerator
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from advanced_alchemy.service import ServiceComposition, ServiceProvider, ServiceWithSession
+
+pytestmark = pytest.mark.unit
+
+
+class FakeServiceA:
+    def __init__(self, *, session: AsyncSession) -> None:
+        self.session = session
+
+
+class FakeServiceB:
+    def __init__(self, *, session: AsyncSession) -> None:
+        self.session = session
+
+
+class FakeServiceC:
+    def __init__(self, *, session: AsyncSession) -> None:
+        self.session = session
+
+
+async def provider_a(session: AsyncSession) -> AsyncGenerator[FakeServiceA, None]:
+    yield FakeServiceA(session=session)
+
+
+async def provider_b(session: AsyncSession) -> AsyncGenerator[FakeServiceB, None]:
+    yield FakeServiceB(session=session)
+
+
+async def provider_c(session: AsyncSession) -> AsyncGenerator[FakeServiceC, None]:
+    yield FakeServiceC(session=session)
+
+
+def test_service_provider_alias_accepts_async_generator_provider() -> None:
+    provider: ServiceProvider[FakeServiceA] = provider_a
+
+    assert provider is provider_a
+
+
+@pytest.mark.asyncio
+async def test_starting_with_returns_typed_composition() -> None:
+    session = MagicMock(spec=AsyncSession)
+    composer = ServiceComposition.starting_with(provider_a)
+
+    async with composer.open(session=session) as services:
+        (a,) = services
+
+    assert isinstance(a, FakeServiceA)
+    assert a.session is session
+
+
+@pytest.mark.asyncio
+async def test_chain_three_providers_share_session() -> None:
+    session = MagicMock(spec=AsyncSession)
+    composer = ServiceComposition.starting_with(provider_a).add(provider_b).add(provider_c)
+
+    async with composer.open(session=session) as services:
+        a, b, c = services
+
+    assert isinstance(a, FakeServiceA)
+    assert isinstance(b, FakeServiceB)
+    assert isinstance(c, FakeServiceC)
+    assert a.session is b.session is c.session is session
+
+
+@pytest.mark.asyncio
+async def test_partial_failure_closes_entered_providers() -> None:
+    cleanup_ran = False
+
+    async def cleanup_provider(session: AsyncSession) -> AsyncGenerator[FakeServiceA, None]:
+        try:
+            yield FakeServiceA(session=session)
+        finally:
+            nonlocal cleanup_ran
+            cleanup_ran = True
+
+    async def failing_provider(session: AsyncSession) -> AsyncGenerator[Any, None]:
+        if repr(session) == "unused":
+            yield None
+        raise RuntimeError("fail at entry")
+
+    session = MagicMock(spec=AsyncSession)
+    composer = ServiceComposition.starting_with(cleanup_provider).add(failing_provider)
+
+    with pytest.raises(RuntimeError, match="fail at entry"):
+        async with composer.open(session=session):
+            pass
+
+    assert cleanup_ran
+
+
+@pytest.mark.asyncio
+async def test_user_exception_inside_block_closes_entered_providers() -> None:
+    cleanup_ran = False
+
+    async def cleanup_provider(session: AsyncSession) -> AsyncGenerator[FakeServiceA, None]:
+        try:
+            yield FakeServiceA(session=session)
+        finally:
+            nonlocal cleanup_ran
+            cleanup_ran = True
+
+    session = MagicMock(spec=AsyncSession)
+
+    async def raise_user_error() -> None:
+        async with ServiceComposition.starting_with(cleanup_provider).open(session=session):
+            raise ValueError("user error")
+
+    with pytest.raises(ValueError, match="user error"):
+        await raise_user_error()
+
+    assert cleanup_ran
+
+
+def test_service_with_session_protocol_recognizes_session_keyword_services() -> None:
+    class GoodService:
+        def __init__(self, *, session: AsyncSession) -> None:
+            self.session = session
+
+    service = GoodService(session=MagicMock(spec=AsyncSession))
+
+    assert isinstance(service, ServiceWithSession)

tests/unit/test_service/test_composition_typing.py

@@ -0,0 +1,79 @@
+"""Static typing assertions for service composition primitives."""
+
+from collections.abc import AsyncGenerator
+from typing import TYPE_CHECKING
+
+from typing_extensions import assert_type
+
+from advanced_alchemy.service import ServiceComposition
+
+if TYPE_CHECKING:
+    from sqlalchemy.ext.asyncio import AsyncSession
+
+
+class S1:
+    pass
+
+
+class S2:
+    pass
+
+
+class S3:
+    pass
+
+
+class S4:
+    pass
+
+
+class S5:
+    pass
+
+
+class S6:
+    pass
+
+
+class S7:
+    pass
+
+
+class S8:
+    pass
+
+
+def p1(session: "AsyncSession") -> AsyncGenerator[S1, None]: ...  # type: ignore[empty-body]
+
+
+def p2(session: "AsyncSession") -> AsyncGenerator[S2, None]: ...  # type: ignore[empty-body]
+
+
+def p3(session: "AsyncSession") -> AsyncGenerator[S3, None]: ...  # type: ignore[empty-body]
+
+
+def p4(session: "AsyncSession") -> AsyncGenerator[S4, None]: ...  # type: ignore[empty-body]
+
+
+def p5(session: "AsyncSession") -> AsyncGenerator[S5, None]: ...  # type: ignore[empty-body]
+
+
+def p6(session: "AsyncSession") -> AsyncGenerator[S6, None]: ...  # type: ignore[empty-body]
+
+
+def p7(session: "AsyncSession") -> AsyncGenerator[S7, None]: ...  # type: ignore[empty-body]
+
+
+def p8(session: "AsyncSession") -> AsyncGenerator[S8, None]: ...  # type: ignore[empty-body]
+
+
+def test_arity_three_typing() -> None:
+    composer = ServiceComposition.starting_with(p1).add(p2).add(p3)
+
+    assert_type(composer, ServiceComposition[S1, S2, S3])
+
+
+def test_arity_eight_typing() -> None:
+    composer = ServiceComposition.starting_with(p1).add(p2).add(p3).add(p4).add(p5).add(p6).add(p7).add(p8)
+
+    assert_type(composer, ServiceComposition[S1, S2, S3, S4, S5, S6, S7, S8])