Create-knowledge-find-decode/query_planning/repositories/postgres.py

"""PostgreSQL adapter for query-planning-owned tables only."""
from __future__ import annotations

from typing import Any
from uuid import UUID

import psycopg2.extras

from query_planning.domain import KnowledgeNeedDraft


Json = psycopg2.extras.Json
psycopg2.extras.register_uuid()


class PostgresQueryPlanningStore:
    def __init__(self, conn: Any) -> None:
        self.conn = conn

    def _id(self, sql: str, params: tuple[Any, ...]) -> UUID:
        with self.conn.cursor() as cur:
            cur.execute(sql, params)
            row = cur.fetchone()
        if row is None:
            raise RuntimeError("query planning insert returned no id")
        return row[0]

    def create_plan(self, **kwargs: Any) -> UUID:
        return self._id(
            """
            INSERT INTO creation_knowledge.query_plans(
                generator_kind, status, input_snapshot, generator_config,
                max_queries, generated_count, normalized_count, unique_count,
                selected_count, dropped_count, metadata
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            RETURNING id
            """,
            (
                kwargs["generator_kind"],
                kwargs["status"],
                Json(kwargs.get("input_snapshot") or {}),
                Json(kwargs.get("generator_config") or {}),
                kwargs.get("max_queries"),
                kwargs["generated_count"],
                kwargs["normalized_count"],
                kwargs["unique_count"],
                kwargs["selected_count"],
                kwargs["dropped_count"],
                Json(kwargs.get("metadata") or {}),
            ),
        )

    def add_knowledge_need(self, *, plan_id: UUID, need: KnowledgeNeedDraft) -> UUID:
        return self._id(
            """
            INSERT INTO creation_knowledge.knowledge_needs(
                plan_id, need_key, decision_context, unknown_information,
                source_ref, priority, metadata
            )
            VALUES (%s, %s, %s, %s, %s, %s, %s)
            RETURNING id
            """,
            (
                plan_id,
                need.need_key,
                need.decision_context,
                need.unknown_information,
                Json(need.source_ref),
                need.priority,
                Json(need.metadata),
            ),
        )

    def link_plan_batch(self, *, plan_id: UUID, batch_id: UUID) -> None:
        with self.conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO creation_knowledge.query_plan_batches(plan_id, batch_id)
                VALUES (%s, %s)
                ON CONFLICT DO NOTHING
                """,
                (plan_id, batch_id),
            )

    def link_query_need(self, *, query_id: UUID, knowledge_need_id: UUID) -> None:
        with self.conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO creation_knowledge.query_knowledge_need_links(
                    query_id, knowledge_need_id
                )
                VALUES (%s, %s)
                ON CONFLICT DO NOTHING
                """,
                (query_id, knowledge_need_id),
            )