python-repositories/python_repositories/adapters/postgres_adapter.py

"""Definition of PostgresAdapter class."""

from __future__ import annotations

from typing import Any

from python_repositories.adapters.connection_aware_adapter import (
    ConnectionAwareAdapter,
)
from python_repositories.config import PostgresConfig
from python_repositories.interfaces import TableRepositoryInterface

try:
    import psycopg
    from psycopg import sql
    from psycopg.errors import UndefinedTable
    from psycopg.rows import dict_row
except ImportError as exc:
    raise ImportError(
        "Postgres support requires the postgres extra. "
        "Install with: pip install python-repositories[postgres]"
    ) from exc


class PostgresAdapter(ConnectionAwareAdapter, TableRepositoryInterface):
    """Postgres adapter exposing basic table CRUD functionality."""

    uri_env_var_name: str = "POSTGRES_URI"
    table_env_var_name: str = "POSTGRES_TABLE"
    primary_key_env_var_name: str = "POSTGRES_PRIMARY_KEY"
    connection_name: str = "Postgres"

    def __init__(
        self,
        *,
        config: PostgresConfig | None = None,
        client: psycopg.Connection[Any] | None = None,
    ) -> None:
        super().__init__()
        if client is not None and config is None:
            raise ValueError("config is required when client is provided")
        if config is None:
            config = PostgresConfig.from_env(
                self.uri_env_var_name,
                self.table_env_var_name,
                self.primary_key_env_var_name,
            )
        self._config = config
        self._client_injected = client is not None
        self._client: psycopg.Connection[Any] | None = client
        if self._client_injected:
            self._invalidate_health_cache()

    def _is_client_ready(self) -> bool:
        return self._client is not None

    def _table_identifier(self) -> sql.Identifier:
        return sql.Identifier(self._config.table)

    def _primary_key_identifier(self) -> sql.Identifier:
        return sql.Identifier(self._config.primary_key)

    def _validate_injected_client(self) -> None:
        if self._client is None:
            return
        try:
            self._run_connection_probe()
        except (psycopg.Error, ConnectionError) as exc:
            raise ConnectionError(
                f"Could not connect to Postgres at {self._config.uri}"
            ) from exc

    def _establish_connection(self) -> None:
        uri = self._config.uri
        try:
            client = psycopg.connect(
                uri,
                row_factory=dict_row,
                connect_timeout=10,
                autocommit=True,
            )
            self._client = client
            self._run_connection_probe()
        except ConnectionError:
            self._client = None
            raise
        except psycopg.Error as exc:
            self._client = None
            raise ConnectionError(f"Could not connect to Postgres at {uri}") from exc

    def _run_connection_probe(self) -> None:
        assert self._client is not None
        with self._client.cursor() as cursor:
            cursor.execute("SELECT 1")
            try:
                cursor.execute(
                    sql.SQL("SELECT 1 FROM {} LIMIT 0").format(self._table_identifier())
                )
            except UndefinedTable as exc:
                raise ConnectionError(
                    f"Table '{self._config.table}' does not exist on Postgres at "
                    f"{self._config.uri}"
                ) from exc

    def disconnect(self) -> None:
        """Disconnect from the Postgres server."""
        if self._client is not None and not self._client_injected:
            self._client.close()
        self._client = None
        self._invalidate_health_cache()

    def _probe_connection(self) -> bool:
        assert self._client is not None
        try:
            self._run_connection_probe()
        except (psycopg.Error, ConnectionError):
            return False
        return True

    def _validate_pk(self, pk: Any) -> None:
        if pk is None:
            raise ValueError("Primary key must not be None")

    def _validate_row(self, row: dict[str, Any]) -> None:
        if not isinstance(row, dict):
            raise ValueError("Row must be a dictionary")
        if self._config.primary_key not in row:
            raise ValueError(
                f"Row must include primary key column '{self._config.primary_key}'"
            )

    def _validate_limit(self, limit: int | None) -> None:
        if limit is not None and (not isinstance(limit, int) or limit < 0):
            raise ValueError("Limit must be a non-negative integer or None")

    def fetch_one(self, pk: Any) -> dict[str, Any] | None:
        """Fetch a single row by primary key."""
        self._validate_pk(pk)
        self._require_connected()
        assert self._client is not None
        query = sql.SQL("SELECT * FROM {} WHERE {} = %s").format(
            self._table_identifier(),
            self._primary_key_identifier(),
        )
        with self._client.cursor() as cursor:
            cursor.execute(query, (pk,))
            row = cursor.fetchone()
        self.logger.debug("Fetched row", pk=pk, found=row is not None)
        return row

    def fetch_all(self, *, limit: int | None = None) -> list[dict[str, Any]]:
        """Fetch all rows from the configured table."""
        self._validate_limit(limit)
        self._require_connected()
        assert self._client is not None
        query = sql.SQL("SELECT * FROM {}").format(self._table_identifier())
        params: tuple[Any, ...] = ()
        if limit is not None:
            query = sql.Composed([query, sql.SQL(" LIMIT %s")])
            params = (limit,)
        with self._client.cursor() as cursor:
            cursor.execute(query, params)
            rows = cursor.fetchall()
        self.logger.debug("Fetched rows", count=len(rows), limit=limit)
        return list(rows)

    def upsert(self, row: dict[str, Any]) -> None:
        """Insert or update a row by primary key."""
        self._validate_row(row)
        self._require_connected()
        assert self._client is not None
        pk_col = self._config.primary_key
        columns = list(row.keys())
        identifiers = [sql.Identifier(column) for column in columns]
        placeholders = sql.SQL(", ").join(sql.Placeholder() * len(columns))
        column_list = sql.SQL(", ").join(identifiers)
        update_columns = [column for column in columns if column != pk_col]
        on_conflict: sql.Composed | sql.SQL
        if update_columns:
            update_assignments = sql.SQL(", ").join(
                sql.SQL("{} = EXCLUDED.{}").format(
                    sql.Identifier(column),
                    sql.Identifier(column),
                )
                for column in update_columns
            )
            on_conflict = sql.SQL("DO UPDATE SET {}").format(update_assignments)
        else:
            on_conflict = sql.SQL("DO NOTHING")
        query = sql.SQL("INSERT INTO {} ({}) VALUES ({}) ON CONFLICT ({}) {}").format(
            self._table_identifier(),
            column_list,
            placeholders,
            self._primary_key_identifier(),
            on_conflict,
        )
        with self._client.cursor() as cursor:
            cursor.execute(query, tuple(row[column] for column in columns))
        self.logger.debug("Upserted row", pk=row[pk_col], columns=columns)

    def delete(self, pk: Any) -> None:
        """Delete a row by primary key."""
        self._validate_pk(pk)
        self._require_connected()
        assert self._client is not None
        query = sql.SQL("DELETE FROM {} WHERE {} = %s").format(
            self._table_identifier(),
            self._primary_key_identifier(),
        )
        with self._client.cursor() as cursor:
            cursor.execute(query, (pk,))
        self.logger.debug("Deleted row", pk=pk)

    def execute(
        self,
        sql_text: str,
        params: tuple[Any, ...] = (),
    ) -> list[dict[str, Any]]:
        """Run a read SQL statement and return rows as dicts."""
        if not isinstance(sql_text, str) or len(sql_text) == 0:
            raise ValueError("SQL must be a non-empty string")
        if not isinstance(params, tuple):
            raise ValueError("Params must be a tuple")
        self._require_connected()
        assert self._client is not None
        with self._client.cursor() as cursor:
            cursor.execute(sql_text, params)
            rows = cursor.fetchall()
        self.logger.debug("Executed SQL", row_count=len(rows))
        return list(rows)