From ab40a4a17e88b4a3d3499c0b4d655d0cc26234d0 Mon Sep 17 00:00:00 2001 From: 1445043649 <> Date: Sun, 28 Jun 2026 11:17:07 +0800 Subject: [PATCH] feat(persistence): ApiKeyRepository with active-key hot path (Stage 1 PR1) Co-Authored-By: Claude Opus 4.8 (1M context) --- .../deerflow/persistence/api_key/__init__.py | 3 +- .../deerflow/persistence/api_key/sql.py | 105 +++++++++++++ backend/tests/test_api_key_repo.py | 140 ++++++++++++++++++ 3 files changed, 247 insertions(+), 1 deletion(-) create mode 100644 backend/packages/harness/deerflow/persistence/api_key/sql.py create mode 100644 backend/tests/test_api_key_repo.py diff --git a/backend/packages/harness/deerflow/persistence/api_key/__init__.py b/backend/packages/harness/deerflow/persistence/api_key/__init__.py index 252b6c7b..b7e76bf2 100644 --- a/backend/packages/harness/deerflow/persistence/api_key/__init__.py +++ b/backend/packages/harness/deerflow/persistence/api_key/__init__.py @@ -14,5 +14,6 @@ middleware live in Stage 1 alongside the headless API surface. from __future__ import annotations from deerflow.persistence.api_key.model import ApiKeyRow +from deerflow.persistence.api_key.sql import ApiKeyRepository -__all__ = ["ApiKeyRow"] +__all__ = ["ApiKeyRepository", "ApiKeyRow"] diff --git a/backend/packages/harness/deerflow/persistence/api_key/sql.py b/backend/packages/harness/deerflow/persistence/api_key/sql.py new file mode 100644 index 00000000..09613de3 --- /dev/null +++ b/backend/packages/harness/deerflow/persistence/api_key/sql.py @@ -0,0 +1,105 @@ +"""SQLAlchemy-backed API key repository (Stage 1 PR1). + +``get_active_by_hash`` is the auth hot path. ``revoked_at IS NULL`` +rides the partial index ``idx_api_keys_active``; expiry is filtered in +Python so the behaviour is identical across sqlite/postgres drivers. + +``_row_to_dict`` deliberately omits ``key_hash`` — no dict this +repository returns ever carries the secret material. +""" + +from __future__ import annotations + +import uuid +from datetime import UTC, datetime +from typing import Any + +from sqlalchemy import select, update +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker + +from deerflow.persistence.api_key.model import ApiKeyRow + + +class ApiKeyRepository: + def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: + self._sf = session_factory + + @staticmethod + def _row_to_dict(row: ApiKeyRow) -> dict[str, Any]: + return { + "id": row.id, + "service_account_id": row.service_account_id, + "key_prefix": row.key_prefix, + "name": row.name, + "scopes": row.scopes, + "rate_limit_rpm": row.rate_limit_rpm, + "expires_at": row.expires_at.isoformat() if row.expires_at else None, + "last_used_at": row.last_used_at.isoformat() if row.last_used_at else None, + "revoked_at": row.revoked_at.isoformat() if row.revoked_at else None, + "created_at": row.created_at.isoformat() if row.created_at else None, + } + + async def create( + self, + *, + service_account_id: str, + key_prefix: str, + key_hash: str, + name: str, + scopes: str, + expires_at: datetime | None = None, + ) -> dict[str, Any]: + row = ApiKeyRow( + id=str(uuid.uuid4()), + service_account_id=service_account_id, + key_prefix=key_prefix, + key_hash=key_hash, + name=name, + scopes=scopes, + expires_at=expires_at, + created_at=datetime.now(UTC), + ) + async with self._sf() as session: + session.add(row) + await session.commit() + await session.refresh(row) + return self._row_to_dict(row) + + async def get(self, key_id: str) -> dict[str, Any] | None: + async with self._sf() as session: + row = await session.get(ApiKeyRow, key_id) + return self._row_to_dict(row) if row else None + + async def get_active_by_hash(self, key_hash: str) -> dict[str, Any] | None: + """Auth hot path: return the key iff not revoked and not expired.""" + async with self._sf() as session: + result = await session.execute(select(ApiKeyRow).where(ApiKeyRow.key_hash == key_hash, ApiKeyRow.revoked_at.is_(None))) + row = result.scalar_one_or_none() + if row is None: + return None + if row.expires_at is not None: + expires_at = row.expires_at + # SQLite (aiosqlite) returns naive datetimes even for DateTime(timezone=True); + # treat them as UTC so the comparison works driver-agnostically. + if expires_at.tzinfo is None: + expires_at = expires_at.replace(tzinfo=UTC) + if expires_at <= datetime.now(UTC): + return None + return self._row_to_dict(row) + + async def list_by_service_account(self, service_account_id: str) -> list[dict[str, Any]]: + async with self._sf() as session: + result = await session.execute(select(ApiKeyRow).where(ApiKeyRow.service_account_id == service_account_id).order_by(ApiKeyRow.created_at.desc())) + return [self._row_to_dict(r) for r in result.scalars()] + + async def revoke(self, key_id: str) -> None: + """Soft-revoke: set ``revoked_at`` (row is kept for audit).""" + async with self._sf() as session: + await session.execute(update(ApiKeyRow).where(ApiKeyRow.id == key_id).values(revoked_at=datetime.now(UTC))) + await session.commit() + + async def touch_last_used(self, key_id: str) -> None: + """Best-effort: stamp ``last_used_at`` after a successful auth.""" + async with self._sf() as session: + await session.execute(update(ApiKeyRow).where(ApiKeyRow.id == key_id).values(last_used_at=datetime.now(UTC))) + await session.commit() diff --git a/backend/tests/test_api_key_repo.py b/backend/tests/test_api_key_repo.py new file mode 100644 index 00000000..e1592ab4 --- /dev/null +++ b/backend/tests/test_api_key_repo.py @@ -0,0 +1,140 @@ +"""Tests for ApiKeyRepository (Stage 1 PR1). + +get_active_by_hash is the auth hot path: must return None for revoked +and expired keys. Expiry is filtered in Python (driver-agnostic) while +revoked_at IS NULL rides the partial index idx_api_keys_active. +""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta + +import pytest + +from deerflow.auth.tokens import generate_api_key +from deerflow.persistence.api_key import ApiKeyRepository +from deerflow.persistence.service_account.model import ServiceAccountRow +from deerflow.persistence.user.model import UserRow +from deerflow.persistence.workspace.model import WorkspaceRow + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend() -> str: + return "asyncio" + + +async def _make_repo(tmp_path): + from deerflow.persistence.engine import get_session_factory, init_engine + + url = f"sqlite+aiosqlite:///{tmp_path / 'test.db'}" + await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path)) + return ApiKeyRepository(get_session_factory()) + + +async def _cleanup(): + from deerflow.persistence.engine import close_engine + + await close_engine() + + +async def _seed_sa(repo, *, sa_id="sa-1") -> None: + async with repo._sf() as session: + session.add(UserRow(id="u-alice", email="alice@example.com")) + await session.commit() + async with repo._sf() as session: + session.add(WorkspaceRow(id="w-1", name="WS", slug="ws", owner_id="u-alice")) + await session.commit() + async with repo._sf() as session: + session.add(ServiceAccountRow(id=sa_id, workspace_id="w-1", name="bot", role="member", identity_mode="collapsed", status="active", created_by="u-alice")) + await session.commit() + + +async def _mint(repo, *, expires_at=None, scopes="threads:read"): + gen = generate_api_key("live") + created = await repo.create(service_account_id="sa-1", key_prefix=gen.prefix, key_hash=gen.key_hash, name="k", scopes=scopes, expires_at=expires_at) + return gen, created + + +async def test_create_then_get_active_by_hash(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + gen, created = await _mint(repo) + assert created["key_prefix"] == gen.prefix + assert "key_hash" not in created # never expose the hash in dicts + found = await repo.get_active_by_hash(gen.key_hash) + assert found is not None + assert found["id"] == created["id"] + assert found["scopes"] == "threads:read" + finally: + await _cleanup() + + +async def test_get_active_by_hash_miss_returns_none(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + assert await repo.get_active_by_hash("deadbeef") is None + finally: + await _cleanup() + + +async def test_revoked_key_not_active(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + gen, created = await _mint(repo) + await repo.revoke(created["id"]) + assert await repo.get_active_by_hash(gen.key_hash) is None + finally: + await _cleanup() + + +async def test_expired_key_not_active(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + past = datetime.now(UTC) - timedelta(hours=1) + gen, _ = await _mint(repo, expires_at=past) + assert await repo.get_active_by_hash(gen.key_hash) is None + finally: + await _cleanup() + + +async def test_future_expiry_still_active(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + future = datetime.now(UTC) + timedelta(hours=1) + gen, _ = await _mint(repo, expires_at=future) + assert await repo.get_active_by_hash(gen.key_hash) is not None + finally: + await _cleanup() + + +async def test_touch_last_used_sets_timestamp(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + gen, created = await _mint(repo) + assert created["last_used_at"] is None + await repo.touch_last_used(created["id"]) + refetched = await repo.get(created["id"]) + assert refetched["last_used_at"] is not None + finally: + await _cleanup() + + +async def test_list_by_service_account(tmp_path): + repo = await _make_repo(tmp_path) + try: + await _seed_sa(repo) + await _mint(repo) + await _mint(repo) + rows = await repo.list_by_service_account("sa-1") + assert len(rows) == 2 + assert all("key_hash" not in r for r in rows) + finally: + await _cleanup()