refactor(api-keys): answer which workspace owns a key, for other domains

docs/domains.md

@@ -62,15 +62,15 @@ since. A module "runs queries" when it imports a query builder (`select`,
 
 | Measure | Count |
 | --- | --- |
-| Service modules | 111, of which 66 sit flat at the top of `services/` |
+| Service modules | 112, of which 66 sit flat at the top of `services/` |
 | Service modules that run queries | 38, plus 2 that only call `session.get` |
 | Route modules | 73 |
 | Route modules that run queries | 17, plus 1 that only calls `session.get` |
 | Route modules that define Pydantic models inline | 40 |
 | Model modules | 19 |
-| Repository modules | 10: a base, `users_repository.py`, 7 under `tenancy/`, and `overview/overview_repository.py` |
-| Service packages per domain | 3: `services/tools/`, which holds the built-in tool registry and no service yet, `services/overview/` and `services/budgets/`. `services/mail/`, `services/routing/` and `services/tenancy/` are older subpackages |
-| Repository packages per domain | 1: `repositories/overview/`. `repositories/tenancy/` is an older subpackage |
+| Repository modules | 11: a base, `users_repository.py`, 7 under `tenancy/`, `overview/overview_repository.py` and `api_keys/api_key_repository.py` |
+| Service packages per domain | 4: `services/tools/`, which holds the built-in tool registry and no service yet, `services/overview/`, `services/budgets/` and `services/api_keys/`. `services/mail/`, `services/routing/` and `services/tenancy/` are older subpackages |
+| Repository packages per domain | 2: `repositories/overview/` and `repositories/api_keys/`. `repositories/tenancy/` is an older subpackage |
 | Modules in `schemas/` | Two domain modules so far, `budgets.py` and `overview.py` |
 | Modules in `exceptions/` | The shared error bases in `_base.py`, which the package root re-exports, and one domain module so far, `budget_exceptions.py`. `services/tenancy/errors.py` holds the rest of the tenancy errors in 1,145 lines |
 
@@ -128,7 +128,8 @@ provisioning, the setup guide, and the gateway's billing users.
 The deployment's and the members' API keys, and which models a key may reach.
 
 - Routes: `keys.py`, `organization_keys.py`
-- Services: `model_access.py`, `bootstrap_service.py`
+- Services: `api_keys/`, `model_access.py`, `bootstrap_service.py`
+- Repositories: `api_keys/`
 - Models: `api_keys.py`
 
 ### budgets

src/gateway/repositories/api_keys/__init__.py

@@ -0,0 +1,3 @@
+from gateway.repositories.api_keys.api_key_repository import ApiKeyRepository
+
+__all__ = ["ApiKeyRepository"]

src/gateway/repositories/api_keys/api_key_repository.py

@@ -0,0 +1,26 @@
+import uuid
+from collections.abc import Sequence
+from typing import Never
+
+from sqlalchemy import select
+
+from gateway.core.unit_of_work import UnitOfWork
+from gateway.models.api_keys import APIKey
+from gateway.repositories.base_repository import BaseRepository
+
+
+class ApiKeyRepository(BaseRepository[APIKey, Never, Never]):
+    """Query API keys in the open block of a Unit of Work."""
+
+    def __init__(self, uow: UnitOfWork) -> None:
+        super().__init__(uow, APIKey)
+
+    async def get_workspace_id_for_key(self, key_id: str) -> uuid.UUID | None:
+        """Return the ID of the workspace that owns a key, or None when no key has that ID."""
+        result = await self.db.execute(select(APIKey.workspace_id).where(APIKey.id == key_id))
+        return result.scalar_one_or_none()
+
+    async def get_key_ids_in_workspaces(self, workspace_ids: Sequence[uuid.UUID]) -> list[str]:
+        """Return the IDs of the keys in these workspaces, in no particular order."""
+        result = await self.db.execute(select(APIKey.id).where(APIKey.workspace_id.in_(workspace_ids)))
+        return list(result.scalars().all())

src/gateway/services/api_keys/__init__.py

@@ -0,0 +1,5 @@
+"""The api-keys domain owns the deployment's and the members' API keys."""
+
+from gateway.services.api_keys._service import ApiKeyService
+
+__all__ = ["ApiKeyService"]

src/gateway/services/api_keys/_service.py

@@ -0,0 +1,22 @@
+import uuid
+from collections.abc import Sequence
+
+from gateway.repositories.api_keys import ApiKeyRepository
+
+
+class ApiKeyService:
+    """Answer questions about API keys and the workspaces that hold them.
+
+    Each method runs in the caller's Unit of Work block and raises ``OutsideUnitOfWorkError`` outside one.
+    """
+
+    def __init__(self, keys: ApiKeyRepository) -> None:
+        self._keys = keys
+
+    async def get_workspace_id_for_key(self, key_id: str) -> uuid.UUID | None:
+        """Return the ID of the workspace that owns a key, or None when no key has that ID."""
+        return await self._keys.get_workspace_id_for_key(key_id)
+
+    async def get_key_ids_in_workspaces(self, workspace_ids: Sequence[uuid.UUID]) -> list[str]:
+        """Return the IDs of the keys in these workspaces, in no particular order."""
+        return await self._keys.get_key_ids_in_workspaces(workspace_ids)

tests/integration/test_api_key_service_lookups.py

@@ -0,0 +1,72 @@
+"""The api-keys service answers which workspace owns a key and which keys sit in a set of workspaces."""
+
+import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from gateway.core.unit_of_work import OutsideUnitOfWorkError, UnitOfWork
+from gateway.models.api_keys import APIKey
+from gateway.models.tenancy import Workspace
+from gateway.repositories.api_keys import ApiKeyRepository
+from gateway.repositories.tenancy import OrganizationRepository, WorkspaceRepository
+from gateway.services.api_keys import ApiKeyService
+
+pytestmark = pytest.mark.asyncio
+
+
+async def _two_workspaces_with_keys(db: AsyncSession) -> tuple[Workspace, Workspace]:
+    organization = await OrganizationRepository(db).create_organization(
+        name="Acme", slug="acme", created_by_user_id=None
+    )
+    workspaces = WorkspaceRepository(db)
+    first = await workspaces.create_workspace(name="First", organization_id=organization.id, created_by_user_id=None)
+    second = await workspaces.create_workspace(name="Second", organization_id=organization.id, created_by_user_id=None)
+    for key_id, workspace in (("sk-first-a", first), ("sk-first-b", first), ("sk-second", second)):
+        db.add(APIKey(id=key_id, key_hash=f"hash-{key_id}", workspace_id=workspace.id))
+    await db.flush()
+    return first, second
+
+
+def _service(uow: UnitOfWork) -> ApiKeyService:
+    return ApiKeyService(ApiKeyRepository(uow))
+
+
+async def test_get_workspace_id_for_key_names_the_owning_workspace(async_db: AsyncSession) -> None:
+    first, second = await _two_workspaces_with_keys(async_db)
+    uow = UnitOfWork(async_db)
+
+    async with uow:
+        assert await _service(uow).get_workspace_id_for_key("sk-first-a") == first.id
+        assert await _service(uow).get_workspace_id_for_key("sk-second") == second.id
+
+
+async def test_get_workspace_id_for_key_is_none_for_an_unknown_key(async_db: AsyncSession) -> None:
+    await _two_workspaces_with_keys(async_db)
+    uow = UnitOfWork(async_db)
+
+    async with uow:
+        assert await _service(uow).get_workspace_id_for_key("sk-unknown") is None
+
+
+async def test_get_key_ids_in_workspaces_returns_only_their_keys(async_db: AsyncSession) -> None:
+    first, second = await _two_workspaces_with_keys(async_db)
+    uow = UnitOfWork(async_db)
+
+    async with uow:
+        service = _service(uow)
+        assert sorted(await service.get_key_ids_in_workspaces([first.id])) == ["sk-first-a", "sk-first-b"]
+        assert sorted(await service.get_key_ids_in_workspaces([first.id, second.id])) == [
+            "sk-first-a",
+            "sk-first-b",
+            "sk-second",
+        ]
+        assert await service.get_key_ids_in_workspaces([]) == []
+
+
+async def test_lookups_raise_outside_a_unit_of_work_block(async_db: AsyncSession) -> None:
+    first, _ = await _two_workspaces_with_keys(async_db)
+    service = _service(UnitOfWork(async_db))
+
+    with pytest.raises(OutsideUnitOfWorkError):
+        await service.get_workspace_id_for_key("sk-first-a")
+    with pytest.raises(OutsideUnitOfWorkError):
+        await service.get_key_ids_in_workspaces([first.id])