"""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)