diff --git a/docs/domains.md b/docs/domains.md index 8212675678..987a2ed3b2 100644 --- a/docs/domains.md +++ b/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 diff --git a/src/gateway/repositories/api_keys/__init__.py b/src/gateway/repositories/api_keys/__init__.py new file mode 100644 index 0000000000..4a76bc52b3 --- /dev/null +++ b/src/gateway/repositories/api_keys/__init__.py @@ -0,0 +1,3 @@ +from gateway.repositories.api_keys.api_key_repository import ApiKeyRepository + +__all__ = ["ApiKeyRepository"] diff --git a/src/gateway/repositories/api_keys/api_key_repository.py b/src/gateway/repositories/api_keys/api_key_repository.py new file mode 100644 index 0000000000..7970d9f112 --- /dev/null +++ b/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()) diff --git a/src/gateway/services/api_keys/__init__.py b/src/gateway/services/api_keys/__init__.py new file mode 100644 index 0000000000..b88fd7a79a --- /dev/null +++ b/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"] diff --git a/src/gateway/services/api_keys/_service.py b/src/gateway/services/api_keys/_service.py new file mode 100644 index 0000000000..1f96e2ddf4 --- /dev/null +++ b/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) diff --git a/tests/integration/test_api_key_service_lookups.py b/tests/integration/test_api_key_service_lookups.py new file mode 100644 index 0000000000..6338a806c9 --- /dev/null +++ b/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])