resume-agent/backend/app/db/repositories.py

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any
from uuid import uuid4

from sqlalchemy import Engine, func, insert, select, update
from sqlalchemy.exc import IntegrityError

from .schema import build_session_tables


class RevisionConflict(Exception):
    """The caller attempted to replace a stale resume document."""


def _now() -> datetime:
    return datetime.now(UTC)


class PostgresSessionRepository:
    """Core PostgreSQL store, shaped for incremental replacement of SQLite."""

    def __init__(self, engine: Engine, *, schema: str = "resume_agent") -> None:
        self.engine = engine
        self.schema = schema
        self.metadata, self.tables = build_session_tables(schema)

    def initialize(self) -> None:
        with self.engine.begin() as connection:
            connection.exec_driver_sql(f'CREATE SCHEMA IF NOT EXISTS "{self.schema}"')
            self.metadata.create_all(connection)

    def create_session(self, session_id: str, stage: str, profile: dict[str, Any]) -> dict[str, Any]:
        now = _now()
        values = {
            "id": session_id,
            "stage": stage,
            "revision": 0,
            "profile": profile,
            "created_at": now,
            "updated_at": now,
        }
        with self.engine.begin() as connection:
            connection.execute(insert(self.tables["sessions"]).values(**values))
        return self.get_session(session_id) or values

    def get_session(self, session_id: str) -> dict[str, Any] | None:
        with self.engine.connect() as connection:
            row = connection.execute(
                select(self.tables["sessions"]).where(self.tables["sessions"].c.id == session_id)
            ).mappings().first()
        return dict(row) if row else None

    def insert_turn(
        self, session_id: str, *, role: str, content: str | None, composer_mode: str, blocks: list[dict[str, Any]]
    ) -> dict[str, Any]:
        turns, block_table = self.tables["turns"], self.tables["blocks"]
        now, turn_id = _now(), f"turn_{uuid4().hex}"
        with self.engine.begin() as connection:
            maximum = connection.execute(
                select(func.coalesce(func.max(turns.c.sequence), 0)).where(turns.c.session_id == session_id)
            ).scalar_one()
            sequence = maximum + 1
            connection.execute(insert(turns).values(
                id=turn_id, session_id=session_id, sequence=sequence, role=role,
                content=content, composer_mode=composer_mode, created_at=now,
            ))
            for index, block in enumerate(blocks):
                connection.execute(insert(block_table).values(
                    id=block.get("id", f"block_{uuid4().hex}"), session_id=session_id,
                    turn_id=turn_id, block_index=index, type=block["type"],
                    lifecycle=block.get("lifecycle", "active"), data=block.get("data", {}),
                    version=1, created_at=now, updated_at=now,
                ))
        return self._turn(turn_id)

    def list_turns(self, session_id: str) -> list[dict[str, Any]]:
        turns = self.tables["turns"]
        with self.engine.connect() as connection:
            ids = connection.execute(
                select(turns.c.id).where(turns.c.session_id == session_id).order_by(turns.c.sequence)
            ).scalars().all()
        return [self._turn(turn_id) for turn_id in ids]

    def _turn(self, turn_id: str) -> dict[str, Any]:
        turns, blocks = self.tables["turns"], self.tables["blocks"]
        with self.engine.connect() as connection:
            turn = connection.execute(select(turns).where(turns.c.id == turn_id)).mappings().one()
            rows = connection.execute(
                select(blocks).where(blocks.c.turn_id == turn_id).order_by(blocks.c.block_index)
            ).mappings().all()
        result = dict(turn)
        result["blocks"] = [dict(row) for row in rows]
        return result

    def create_resume(
        self, session_id: str, resume_id: str, idempotency_key: str | None, content: dict[str, Any]
    ) -> dict[str, Any]:
        existing = self.get_resume(session_id)
        if existing:
            return existing
        now = _now()
        try:
            with self.engine.begin() as connection:
                connection.execute(insert(self.tables["resumes"]).values(
                    id=resume_id, session_id=session_id, idempotency_key=idempotency_key,
                    revision=1, content=content, created_at=now, updated_at=now,
                ))
        except IntegrityError:
            existing = self.get_resume(session_id)
            if existing:
                return existing
            raise
        return self.get_resume(session_id)  # type: ignore[return-value]

    def get_resume(self, session_id: str) -> dict[str, Any] | None:
        with self.engine.connect() as connection:
            row = connection.execute(
                select(self.tables["resumes"]).where(self.tables["resumes"].c.session_id == session_id)
            ).mappings().first()
        return dict(row) if row else None

    def update_resume(self, session_id: str, content: dict[str, Any], *, expected_revision: int) -> dict[str, Any]:
        resumes = self.tables["resumes"]
        with self.engine.begin() as connection:
            result = connection.execute(update(resumes).where(
                resumes.c.session_id == session_id, resumes.c.revision == expected_revision
            ).values(content=content, revision=expected_revision + 1, updated_at=_now()))
        if result.rowcount != 1:
            raise RevisionConflict(session_id)
        return self.get_resume(session_id)  # type: ignore[return-value]

    def create_resume_import(
        self,
        session_id: str,
        *,
        import_id: str,
        file_name: str,
        mime_type: str,
        size_bytes: int,
        sha256: str,
        object_key: str,
        document: dict[str, Any] | None,
        field_reviews: list[dict[str, Any]],
        status: str = "awaiting_review",
        error_code: str | None = None,
    ) -> dict[str, Any]:
        existing = self.find_resume_import_by_sha256(session_id, sha256)
        if existing:
            return existing
        now = _now()
        values = {
            "id": import_id,
            "session_id": session_id,
            "file_name": file_name,
            "mime_type": mime_type,
            "size_bytes": size_bytes,
            "sha256": sha256,
            "object_key": object_key,
            "status": status,
            "document": document,
            "field_reviews": field_reviews,
            "error_code": error_code,
            "created_at": now,
            "updated_at": now,
        }
        try:
            with self.engine.begin() as connection:
                connection.execute(insert(self.tables["resume_imports"]).values(**values))
        except IntegrityError:
            existing = self.find_resume_import_by_sha256(session_id, sha256)
            if existing:
                return existing
            raise
        return self.get_resume_import(session_id, import_id)  # type: ignore[return-value]

    def get_resume_import(self, session_id: str, import_id: str) -> dict[str, Any] | None:
        imports = self.tables["resume_imports"]
        with self.engine.connect() as connection:
            row = connection.execute(
                select(imports).where(
                    imports.c.id == import_id,
                    imports.c.session_id == session_id,
                )
            ).mappings().first()
        return dict(row) if row else None

    def find_resume_import_by_sha256(self, session_id: str, sha256: str) -> dict[str, Any] | None:
        imports = self.tables["resume_imports"]
        with self.engine.connect() as connection:
            row = connection.execute(
                select(imports).where(
                    imports.c.session_id == session_id,
                    imports.c.sha256 == sha256,
                )
            ).mappings().first()
        return dict(row) if row else None

    def update_resume_import_status(
        self,
        session_id: str,
        import_id: str,
        status: str,
        *,
        error_code: str | None = None,
    ) -> dict[str, Any]:
        imports = self.tables["resume_imports"]
        with self.engine.begin() as connection:
            result = connection.execute(
                update(imports)
                .where(imports.c.id == import_id, imports.c.session_id == session_id)
                .values(status=status, error_code=error_code, updated_at=_now())
            )
        if result.rowcount != 1:
            raise KeyError(import_id)
        return self.get_resume_import(session_id, import_id)  # type: ignore[return-value]
    def delete_session(self, session_id: str) -> bool:
        with self.engine.begin() as connection:
            result = connection.execute(
                self.tables["sessions"].delete().where(self.tables["sessions"].c.id == session_id)
            )
        return result.rowcount == 1