feat(persistence): ApiKeyRepository with active-key hot path (Stage 1 PR1)
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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"]
|
||||
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user