"""PostgreSQL, nur lesend (Grundlage: Konnektor von Bifröst). Sicherungen: - Verbindung mit default_transaction_read_only; jede Abfrage in einer schreibgeschützten Transaktion, die immer zurückgerollt wird - Ausführung über einen serverseitigen Cursor (DECLARE … CURSOR), der keine zweite Anweisung zulässt - Zeitlimit in der Datenbank (statement_timeout) und zusätzlich vom Client aus - höchstens max_rows Zeilen werden geholt """ import threading from collections.abc import Iterator from contextlib import contextmanager from typing import Any import psycopg from psycopg.conninfo import make_conninfo from dbmcp.core import Invalid from dbmcp.core.config import Settings from dbmcp.core.connection import Connection from dbmcp.core.driver import Check from dbmcp.core.result import Table SYSTEM = ("pg_catalog", "information_schema") # Bekannte Ursachen beim Verbindungsaufbau; libpq nennt bei mehreren Adressen alle Versuche CAUSES = ( ("server does not support SSL", "Der Server unterstützt kein TLS, TLS ist aber erzwungen."), ("password authentication failed", "Benutzer oder Passwort falsch."), ("no password supplied", "Der Server verlangt ein Passwort."), ("does not exist", None), ("no pg_hba.conf entry", "Der Server lässt diesen Zugang nicht zu (pg_hba.conf)."), ("timeout expired", "Zeitüberschreitung beim Verbindungsaufbau; Host oder Port prüfen."), ("Connection refused", "Verbindung abgelehnt; läuft der Server und stimmt der Port?"), ) def first_line(error: Exception) -> str: text = str(error).strip() return text.splitlines()[0][:300] if text else type(error).__name__ def connect_error(error: Exception) -> str: text = str(error) for needle, message in CAUSES: if needle in text: if message is None: # z. B. database "x" does not exist: Originaltext ist verständlich line = next(l for l in text.splitlines() if needle in l) return line.rsplit("failed: ", 1)[-1].strip()[:300] return message return first_line(error) def quote(identifier: str) -> str: return '"' + identifier.replace('"', '""') + '"' class PostgresSession: def __init__(self, con: psycopg.Connection, settings: Settings): self.con, self.settings = con, settings def fetch(self, sql: str, params: tuple[Any, ...] | None, max_rows: int) -> Table: # Zweite Sicherung, falls die Abfrage das Zeitlimit in der Datenbank aushebelt timer = threading.Timer(self.settings.query_timeout + 2, self.con.cancel_safe) timer.start() try: with self.con.cursor(name="dbmcp_query") as cursor: cursor.execute(sql, params) columns = [d.name for d in cursor.description or []] rows = cursor.fetchmany(max_rows + 1) except psycopg.errors.QueryCanceled as e: raise Invalid(f"Abfrage nach {self.settings.query_timeout:g} Sekunden abgebrochen; bitte eingrenzen " "oder aggregieren.") from e except psycopg.Error as e: raise Invalid(f"SQL-Fehler: {first_line(e)}") from e finally: timer.cancel() return Table(columns, [list(r) for r in rows[:max_rows]], len(rows) > max_rows) def check(self) -> Check: try: version = self.con.execute("SHOW server_version").fetchone()[0] tables = self.con.execute( "SELECT count(*) FROM information_schema.tables WHERE table_schema NOT IN (%s, %s)", SYSTEM ).fetchone()[0] writable = self.con.execute( "SELECT count(DISTINCT (table_schema, table_name)) FROM information_schema.table_privileges WHERE grantee = current_user " "AND privilege_type IN ('INSERT', 'UPDATE', 'DELETE', 'TRUNCATE') AND table_schema NOT IN (%s, %s)", SYSTEM ).fetchone()[0] except psycopg.Error as e: raise Invalid(f"Fehler: {first_line(e)}") from e return Check(version, tables, writable) def tables(self, limit: int) -> Table: table = self.fetch( "SELECT t.table_schema, t.table_name, CASE WHEN t.table_type = 'VIEW' THEN 'view' ELSE 'table' END, " "(SELECT c.reltuples::bigint FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace " " WHERE n.nspname = t.table_schema AND c.relname = t.table_name), " "obj_description(format('%%I.%%I', t.table_schema, t.table_name)::regclass) " "FROM information_schema.tables t WHERE t.table_schema NOT IN (%s, %s) " "ORDER BY t.table_schema, t.table_name", SYSTEM, limit) table.columns = ["schema", "name", "type", "rows_estimate", "comment"] return table def describe(self, table: str, sample_rows: int) -> tuple[Table, Table]: schema, _, name = table.rpartition(".") schema, name = schema.strip('"') or "public", name.strip('"') columns = self.fetch( "SELECT column_name, data_type, is_nullable = 'YES', " "col_description(format('%%I.%%I', table_schema, table_name)::regclass, ordinal_position) " "FROM information_schema.columns WHERE table_schema = %s AND table_name = %s ORDER BY ordinal_position", (schema, name), 1000) if not columns.rows: raise Invalid(f"Tabelle {schema}.{name} gibt es nicht oder sie ist nicht sichtbar; siehe list_tables.") columns.columns = ["name", "type", "nullable", "comment"] sample = self.fetch(f"SELECT * FROM {quote(schema)}.{quote(name)} LIMIT {int(sample_rows)}".replace("%", "%%"), (), sample_rows) return columns, sample def query(self, sql: str, max_rows: int) -> Table: return self.fetch(sql, None, max_rows) class PostgresDriver: name = "PostgreSQL" key = "postgres" default_port = 5432 dialect = "PostgreSQL; Schema public, Begrenzen mit LIMIT." starts = ("select", "with", "table", "values") def __init__(self, settings: Settings): self.settings = settings def conninfo(self, connection: Connection) -> str: return make_conninfo(host=connection.host, port=connection.port, dbname=connection.database, user=connection.user, password=connection.password or None, sslmode=connection.sslmode, connect_timeout=max(int(self.settings.connect_timeout), 1), application_name="db-mcp-postgres", options="-c default_transaction_read_only=on") @contextmanager def session(self, connection: Connection) -> Iterator[PostgresSession]: try: con = psycopg.connect(self.conninfo(connection)) except psycopg.Error as e: raise Invalid(f"Datenbank nicht erreichbar: {connect_error(e)}") from e try: con.execute("SET TRANSACTION READ ONLY") con.execute(f"SET LOCAL statement_timeout = {int(self.settings.query_timeout * 1000)}") con.execute("SET LOCAL lock_timeout = 2000") yield PostgresSession(con, self.settings) finally: con.rollback() con.close()