evalkit-py 0.2.0.dev12__py3-none-any.whl

evalkit/db/trace_imports.py

@@ -0,0 +1,308 @@
+from datetime import UTC, datetime, timedelta
+
+from pydantic import BaseModel
+from sqlalchemy import func, select, update
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from evalkit.contracts.enums import TraceImportStatus
+from evalkit.core.traces.importer import (
+    ACTIVE_IMPORT,
+    IMPORT_LEASE_SECONDS,
+    IMPORT_STOPPED_ERROR,
+    ImportJob,
+)
+from evalkit.db import models
+from evalkit.db.locks import lock_name
+from evalkit.store.base import Conflict
+
+ACTIVE = [s.value for s in ACTIVE_IMPORT]
+RESTARTABLE = [TraceImportStatus.FAILED.value, TraceImportStatus.STOPPED.value]
+
+
+class UnknownTraceImport(KeyError):
+    def __init__(self, import_id: int):
+        super().__init__(f"unknown trace import: {import_id}")
+
+
+class ImportRunning(Conflict):
+    def __init__(self, import_id: int):
+        super().__init__(f"import {import_id} is already running for this agent")
+
+
+class ImportNotActive(Conflict):
+    def __init__(self, import_id: int, status: str, action: str):
+        super().__init__(f"import {import_id} is {status}; only {action}")
+
+
+class TraceImportView(BaseModel):
+    id: int
+    agent: str
+    source: str
+    project: str
+    since: datetime
+    until: datetime
+    resume_from: datetime
+    status: TraceImportStatus
+    traces: int
+    spans: int
+    sample_session: str | None
+    error: str | None
+    created_by: str
+    created_at: datetime
+    finished_at: datetime | None
+
+
+def _lease() -> datetime:
+    return datetime.now(UTC) + timedelta(seconds=IMPORT_LEASE_SECONDS)
+
+
+def _views():
+    return (
+        select(
+            models.TraceImport,
+            models.Agent.name,
+            models.Workspace.name,
+            models.TraceSource.name,
+            models.User.username,
+        )
+        .join(models.Agent, models.Agent.id == models.TraceImport.agent_id)
+        .join(models.Workspace, models.Workspace.id == models.Agent.workspace_id)
+        .join(models.TraceSource, models.TraceSource.id == models.TraceImport.trace_source_id)
+        .join(models.User, models.User.id == models.TraceImport.created_by)
+    )
+
+
+def _view(
+    row: models.TraceImport, agent: str, workspace: str, source: str, user: str
+) -> TraceImportView:
+    return TraceImportView(
+        id=row.id,
+        agent=f"{agent}@{workspace}",
+        source=source,
+        project=row.project,
+        since=row.since,
+        until=row.until,
+        resume_from=row.resume_from,
+        status=TraceImportStatus(row.status),
+        traces=row.traces,
+        spans=row.spans,
+        sample_session=row.sample_session,
+        error=row.error,
+        created_by=user,
+        created_at=row.created_at,
+        finished_at=row.finished_at,
+    )
+
+
+async def _refuse_second(session: AsyncSession, agent_id: int) -> None:
+    await lock_name(session, "trace-import", str(agent_id))
+    running = await session.scalar(
+        select(models.TraceImport.id).where(
+            models.TraceImport.agent_id == agent_id, models.TraceImport.status.in_(ACTIVE)
+        )
+    )
+    if running is not None:
+        raise ImportRunning(running)
+
+
+class TraceImportStore:
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
+        self._sessions = session_factory
+
+    async def _one(self, session: AsyncSession, agent_id: int, import_id: int) -> TraceImportView:
+        found = (
+            await session.execute(
+                _views().where(
+                    models.TraceImport.id == import_id, models.TraceImport.agent_id == agent_id
+                )
+            )
+        ).one_or_none()
+        if found is None:
+            raise UnknownTraceImport(import_id)
+        return _view(*found)
+
+    async def create(
+        self,
+        agent_id: int,
+        trace_source_id: int,
+        project: str,
+        since: datetime,
+        until: datetime,
+        created_by: int,
+        holder: str,
+    ) -> TraceImportView:
+        async with self._sessions() as session:
+            await _refuse_second(session, agent_id)
+            row = models.TraceImport(
+                agent_id=agent_id,
+                trace_source_id=trace_source_id,
+                project=project,
+                since=since,
+                until=until,
+                resume_from=since,
+                resume_trace_ids=[],
+                status=TraceImportStatus.RUNNING.value,
+                traces=0,
+                spans=0,
+                lease_holder=holder,
+                lease_expires_at=_lease(),
+                created_by=created_by,
+            )
+            session.add(row)
+            await session.commit()
+            return await self._one(session, agent_id, row.id)
+
+    async def get(self, agent_id: int, import_id: int) -> TraceImportView:
+        async with self._sessions() as session:
+            return await self._one(session, agent_id, import_id)
+
+    async def recent(self, agent_id: int, limit: int = 20) -> list[TraceImportView]:
+        async with self._sessions() as session:
+            rows = await session.execute(
+                _views()
+                .where(models.TraceImport.agent_id == agent_id)
+                .order_by(models.TraceImport.id.desc())
+                .limit(limit)
+            )
+            return [_view(*row) for row in rows.all()]
+
+    async def stop(self, agent_id: int, import_id: int) -> TraceImportView:
+        async with self._sessions() as session:
+            current = await self._one(session, agent_id, import_id)
+            if current.status is not TraceImportStatus.RUNNING:
+                raise ImportNotActive(import_id, current.status, "a running import can be stopped")
+            await session.execute(
+                update(models.TraceImport)
+                .where(
+                    models.TraceImport.id == import_id,
+                    models.TraceImport.status == TraceImportStatus.RUNNING.value,
+                )
+                .values(status=TraceImportStatus.STOPPING.value)
+            )
+            await session.commit()
+            return await self._one(session, agent_id, import_id)
+
+    async def restart(self, agent_id: int, import_id: int, holder: str) -> TraceImportView:
+        async with self._sessions() as session:
+            current = await self._one(session, agent_id, import_id)
+            if current.status not in RESTARTABLE:
+                raise ImportNotActive(
+                    import_id, current.status, "a failed or stopped import can be continued"
+                )
+            await _refuse_second(session, agent_id)
+            await session.execute(
+                update(models.TraceImport)
+                .where(models.TraceImport.id == import_id)
+                .values(
+                    status=TraceImportStatus.RUNNING.value,
+                    error=None,
+                    finished_at=None,
+                    lease_holder=holder,
+                    lease_expires_at=_lease(),
+                )
+            )
+            await session.commit()
+            return await self._one(session, agent_id, import_id)
+
+    async def job(self, import_id: int) -> ImportJob:
+        async with self._sessions() as session:
+            row = await session.get(models.TraceImport, import_id)
+            if row is None:
+                raise UnknownTraceImport(import_id)
+            return ImportJob(
+                id=row.id,
+                agent_id=row.agent_id,
+                trace_source_id=row.trace_source_id,
+                project=row.project,
+                until=row.until,
+                resume_from=row.resume_from,
+                resume_trace_ids=list(row.resume_trace_ids),
+            )
+
+    async def advance(
+        self,
+        import_id: int,
+        holder: str,
+        *,
+        resume_from: datetime,
+        resume_trace_ids: list[str],
+        traces: int,
+        spans: int,
+        sample_session: str | None,
+    ) -> TraceImportStatus | None:
+        async with self._sessions() as session:
+            status = await session.scalar(
+                update(models.TraceImport)
+                .where(
+                    models.TraceImport.id == import_id,
+                    models.TraceImport.lease_holder == holder,
+                    models.TraceImport.status.in_(ACTIVE),
+                )
+                .values(
+                    resume_from=resume_from,
+                    resume_trace_ids=resume_trace_ids,
+                    traces=models.TraceImport.traces + traces,
+                    spans=models.TraceImport.spans + spans,
+                    lease_expires_at=_lease(),
+                    **({"sample_session": sample_session} if sample_session else {}),
+                )
+                .returning(models.TraceImport.status)
+            )
+            await session.commit()
+        return TraceImportStatus(status) if status is not None else None
+
+    async def renew(self, import_id: int, holder: str) -> TraceImportStatus | None:
+        async with self._sessions() as session:
+            status = await session.scalar(
+                update(models.TraceImport)
+                .where(
+                    models.TraceImport.id == import_id,
+                    models.TraceImport.lease_holder == holder,
+                    models.TraceImport.status.in_(ACTIVE),
+                )
+                .values(lease_expires_at=_lease())
+                .returning(models.TraceImport.status)
+            )
+            await session.commit()
+        return TraceImportStatus(status) if status is not None else None
+
+    async def finish(
+        self, import_id: int, holder: str, status: TraceImportStatus, error: str | None = None
+    ) -> None:
+        async with self._sessions() as session:
+            await session.execute(
+                update(models.TraceImport)
+                .where(
+                    models.TraceImport.id == import_id, models.TraceImport.lease_holder == holder
+                )
+                .values(
+                    status=status.value,
+                    error=error,
+                    finished_at=func.now(),
+                    lease_holder=None,
+                    lease_expires_at=None,
+                )
+            )
+            await session.commit()
+
+
+async def reap_imports(session_factory: async_sessionmaker[AsyncSession]) -> list[int]:
+    async with session_factory() as session:
+        reaped = await session.scalars(
+            update(models.TraceImport)
+            .where(
+                models.TraceImport.status.in_(ACTIVE),
+                models.TraceImport.lease_expires_at < datetime.now(UTC),
+            )
+            .values(
+                status=TraceImportStatus.FAILED.value,
+                error=IMPORT_STOPPED_ERROR,
+                finished_at=func.now(),
+                lease_holder=None,
+                lease_expires_at=None,
+            )
+            .returning(models.TraceImport.id)
+        )
+        ids = list(reaped.all())
+        await session.commit()
+    return ids

evalkit/db/trace_sources.py

@@ -0,0 +1,131 @@
+from datetime import UTC, datetime
+
+from pydantic import BaseModel, SecretStr
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from evalkit.contracts.enums import TraceProvider
+from evalkit.db import models
+from evalkit.db.crypto import SecretBox
+from evalkit.db.locks import lock_name
+
+
+class UnknownTraceSource(KeyError):
+    def __init__(self, name: str):
+        super().__init__(f"unknown trace source: {name}")
+
+
+class TraceSourceExists(ValueError):
+    def __init__(self, name: str):
+        super().__init__(f"trace source exists: {name}")
+
+
+class EncryptionKeyMissing(RuntimeError):
+    def __init__(self) -> None:
+        super().__init__("encryption_key is not configured")
+
+
+class TraceSourceView(BaseModel):
+    id: int
+    name: str
+    provider: TraceProvider
+    url: str
+    archived: bool
+
+
+class ResolvedTraceSource(BaseModel):
+    id: int
+    name: str
+    provider: TraceProvider
+    url: str
+    api_key: SecretStr | None
+
+
+def _view(row: models.TraceSource) -> TraceSourceView:
+    return TraceSourceView(
+        id=row.id,
+        name=row.name,
+        provider=TraceProvider(row.provider),
+        url=row.config["url"],
+        archived=row.archived_at is not None,
+    )
+
+
+class TraceSourceStore:
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession], box: SecretBox | None):
+        self._sessions = session_factory
+        self._box = box
+
+    def _seal(self, api_key: str) -> bytes:
+        if self._box is None:
+            raise EncryptionKeyMissing()
+        return self._box.seal(api_key)
+
+    async def _row(self, session: AsyncSession, name: str) -> models.TraceSource:
+        row = await session.scalar(
+            select(models.TraceSource).where(models.TraceSource.name == name)
+        )
+        if row is None:
+            raise UnknownTraceSource(name)
+        return row
+
+    async def create(
+        self, name: str, provider: TraceProvider, url: str, api_key: str | None
+    ) -> TraceSourceView:
+        sealed = self._seal(api_key) if api_key else None
+        async with self._sessions() as session:
+            await lock_name(session, "trace-source", name)
+            if await session.scalar(
+                select(models.TraceSource.id).where(models.TraceSource.name == name)
+            ):
+                raise TraceSourceExists(name)
+            row = models.TraceSource(
+                name=name, provider=provider.value, config={"url": url}, api_key_encrypted=sealed
+            )
+            session.add(row)
+            await session.commit()
+            return _view(row)
+
+    async def edit(
+        self, name: str, *, url: str | None = None, api_key: str | None = None
+    ) -> TraceSourceView:
+        async with self._sessions() as session:
+            row = await self._row(session, name)
+            if url is not None:
+                row.config = {**row.config, "url": url}
+            if api_key is not None:
+                row.api_key_encrypted = self._seal(api_key)
+            await session.commit()
+            return _view(row)
+
+    async def archive(self, name: str) -> TraceSourceView:
+        async with self._sessions() as session:
+            row = await self._row(session, name)
+            row.archived_at = datetime.now(UTC)
+            await session.commit()
+            return _view(row)
+
+    async def list_sources(self, include_archived: bool = False) -> list[TraceSourceView]:
+        query = select(models.TraceSource).order_by(models.TraceSource.name)
+        if not include_archived:
+            query = query.where(models.TraceSource.archived_at.is_(None))
+        async with self._sessions() as session:
+            return [_view(row) for row in (await session.scalars(query)).all()]
+
+    async def resolve(self, name: str) -> ResolvedTraceSource:
+        async with self._sessions() as session:
+            row = await self._row(session, name)
+        if row.archived_at is not None:
+            raise UnknownTraceSource(name)
+        api_key = None
+        if row.api_key_encrypted is not None:
+            if self._box is None:
+                raise EncryptionKeyMissing()
+            api_key = self._box.open(row.api_key_encrypted)
+        return ResolvedTraceSource(
+            id=row.id,
+            name=row.name,
+            provider=TraceProvider(row.provider),
+            url=row.config["url"],
+            api_key=SecretStr(api_key) if api_key is not None else None,
+        )

evalkit/db/users.py

@@ -0,0 +1,186 @@
+import re
+from datetime import UTC, datetime
+
+import bcrypt
+from sqlalchemy import delete, func, select, union_all, update
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
+
+from evalkit.contracts.enums import TokenKind
+from evalkit.db.locks import lock_name
+from evalkit.db.models import AuthToken, User, Workspace
+
+NAME_PATTERN = re.compile(r"^[a-z0-9][a-z0-9._-]*$")
+
+
+class NameTaken(Exception):
+    def __init__(self, name: str):
+        super().__init__(f"name {name!r} is already a user or workspace")
+        self.name = name
+
+
+class InvalidName(ValueError):
+    def __init__(self, name: str):
+        super().__init__(
+            f"invalid name {name!r}: use lowercase letters, digits, '.', '_' or '-', "
+            "starting with a letter or digit"
+        )
+
+
+def check_name(name: str) -> None:
+    if not NAME_PATTERN.fullmatch(name):
+        raise InvalidName(name)
+
+
+async def lock_names(session: AsyncSession) -> None:
+    await lock_name(session, "users-and-workspaces")
+
+
+class LastAdmin(Exception):
+    def __init__(self) -> None:
+        super().__init__("cannot remove the last admin")
+
+
+class UnknownUser(KeyError):
+    def __init__(self, username: str):
+        super().__init__(f"unknown user: {username}")
+
+
+class InvalidPassword(ValueError):
+    pass
+
+
+BCRYPT_ROUNDS = 12
+PASSWORD_MIN_CHARS = 8
+PASSWORD_MAX_BYTES = 72
+
+
+def check_password_size(password: str) -> None:
+    if len(password.encode()) > PASSWORD_MAX_BYTES:
+        raise InvalidPassword(f"password is longer than {PASSWORD_MAX_BYTES} bytes")
+
+
+def check_new_password(password: str) -> None:
+    if len(password) < PASSWORD_MIN_CHARS:
+        raise InvalidPassword(f"password must be at least {PASSWORD_MIN_CHARS} characters")
+    check_password_size(password)
+
+
+def hash_password(password: str) -> str:
+    return bcrypt.hashpw(password.encode(), bcrypt.gensalt(BCRYPT_ROUNDS)).decode()
+
+
+def verify_password(password: str, password_hash: str) -> bool:
+    return bcrypt.checkpw(password.encode(), password_hash.encode())
+
+
+async def name_taken(session: AsyncSession, name: str) -> bool:
+    count = await session.scalar(
+        select(func.count()).select_from(
+            union_all(
+                select(User.id).where(User.username == name),
+                select(Workspace.id).where(Workspace.name == name),
+            ).subquery()
+        )
+    )
+    return bool(count)
+
+
+async def user_id(session: AsyncSession, username: str) -> int:
+    found = await session.scalar(select(User.id).where(User.username == username))
+    if found is None:
+        raise UnknownUser(username)
+    return found
+
+
+class UserStore:
+    def __init__(self, session_factory: async_sessionmaker[AsyncSession]):
+        self._sf = session_factory
+
+    async def create(self, username: str, password: str, *, is_admin: bool = False) -> User:
+        check_name(username)
+        check_new_password(password)
+        async with self._sf() as s:
+            await lock_names(s)
+            if await name_taken(s, username):
+                raise NameTaken(username)
+            user = User(username=username, password_hash=hash_password(password), is_admin=is_admin)
+            s.add(user)
+            await s.flush()
+            s.add(Workspace(name=username, owner_user_id=user.id))
+            await s.commit()
+            return user
+
+    async def _get(self, s: AsyncSession, username: str) -> User:
+        user = await s.scalar(select(User).where(User.username == username))
+        if user is None:
+            raise UnknownUser(username)
+        return user
+
+    async def _ensure_other_admin(self, s: AsyncSession, user: User) -> None:
+        if not user.is_admin or user.archived_at is not None:
+            return
+        await lock_name(s, "admins")
+        others = await s.scalar(
+            select(func.count()).where(
+                User.is_admin.is_(True), User.archived_at.is_(None), User.id != user.id
+            )
+        )
+        if not others:
+            raise LastAdmin()
+
+    async def archive(self, username: str) -> None:
+        now = datetime.now(UTC)
+        async with self._sf() as s:
+            user = await self._get(s, username)
+            await self._ensure_other_admin(s, user)
+            user.archived_at = now
+            await s.execute(
+                delete(AuthToken).where(
+                    AuthToken.user_id == user.id, AuthToken.kind != TokenKind.AGENT.value
+                )
+            )
+            await s.execute(
+                update(Workspace)
+                .where(Workspace.owner_user_id == user.id, Workspace.archived_at.is_(None))
+                .values(archived_at=now)
+            )
+            await s.commit()
+
+    async def restore(self, username: str) -> None:
+        async with self._sf() as s:
+            user = await self._get(s, username)
+            if user.archived_at is not None:
+                await s.execute(
+                    update(Workspace)
+                    .where(
+                        Workspace.owner_user_id == user.id,
+                        Workspace.archived_at == user.archived_at,
+                    )
+                    .values(archived_at=None)
+                )
+            user.archived_at = None
+            await s.commit()
+
+    async def set_password(self, username: str, password: str) -> None:
+        check_new_password(password)
+        async with self._sf() as s:
+            user = await self._get(s, username)
+            user.password_hash = hash_password(password)
+            await s.execute(
+                delete(AuthToken).where(
+                    AuthToken.user_id == user.id, AuthToken.kind != TokenKind.AGENT.value
+                )
+            )
+            await s.commit()
+
+    async def set_admin(self, username: str, value: bool) -> None:
+        async with self._sf() as s:
+            user = await self._get(s, username)
+            if not value:
+                await self._ensure_other_admin(s, user)
+            user.is_admin = value
+            await s.commit()
+
+    async def list(self) -> list[User]:
+        async with self._sf() as s:
+            return list((await s.scalars(select(User).order_by(User.username))).all())