155 lines
7.1 KiB
Python
155 lines
7.1 KiB
Python
"""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()
|