nngoandClaude Sonnet 4.6
4ccc81955e
SSL was wired through the UI and ConnectionProfile but never actually
applied in any driver or passed from the connection builder.
- main_window: pass ssl/ssl_ca/ssl_cert/ssl_key from profile into
the driver config dict so drivers can act on them
- MySQL: add ssl={ca,cert,key} to pymysql connect kwargs when enabled
- PostgreSQL: set sslmode=verify-ca (with CA) or require, plus
sslrootcert/sslcert/sslkey when provided
- MSSQL: switch Encrypt=yes + TrustServerCertificate=no when SSL is
on; Encrypt=no + TrustServerCertificate=yes otherwise
- SQLite: no change needed (local file, no network layer)
WHERE filter: add tooltip explicitly labelling the input as raw SQL
so users understand arbitrary expressions are executed directly.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
349 lines
14 KiB
Python
349 lines
14 KiB
Python
"""PostgreSQL driver implementation using psycopg2."""
|
|
import psycopg2
|
|
import psycopg2.extras
|
|
from contextlib import contextmanager
|
|
from typing import Optional
|
|
from app.drivers.base import (
|
|
BaseDriver, ColumnInfo, IndexInfo, ForeignKeyInfo, TableInfo
|
|
)
|
|
|
|
|
|
class PostgreSQLDriver(BaseDriver):
|
|
"""PostgreSQL database driver."""
|
|
|
|
def __init__(self, config: dict):
|
|
super().__init__(config)
|
|
self.db_type = "postgresql"
|
|
|
|
def _dsn(self) -> dict:
|
|
kw = {
|
|
"host": self.config.get("host", "localhost"),
|
|
"port": int(self.config.get("port", 5432)),
|
|
"user": self.config.get("user", ""),
|
|
"password": self.config.get("password", ""),
|
|
"connect_timeout": int(self.config.get("connection_timeout", 30)),
|
|
}
|
|
db = self.config.get("database", "")
|
|
if db:
|
|
kw["dbname"] = db
|
|
if self.config.get("ssl"):
|
|
kw["sslmode"] = "verify-ca" if self.config.get("ssl_ca") else "require"
|
|
if self.config.get("ssl_ca"): kw["sslrootcert"] = self.config["ssl_ca"]
|
|
if self.config.get("ssl_cert"): kw["sslcert"] = self.config["ssl_cert"]
|
|
if self.config.get("ssl_key"): kw["sslkey"] = self.config["ssl_key"]
|
|
return kw
|
|
|
|
def connect(self) -> None:
|
|
self._connection = psycopg2.connect(**self._dsn())
|
|
self._connection.autocommit = True
|
|
|
|
def disconnect(self) -> None:
|
|
if self._connection:
|
|
try:
|
|
self._connection.close()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
self._connection = None
|
|
|
|
def test_connection(self) -> tuple:
|
|
try:
|
|
conn = psycopg2.connect(**self._dsn())
|
|
conn.close()
|
|
return True, "Connection successful"
|
|
except Exception as e:
|
|
return False, str(e)
|
|
|
|
@contextmanager
|
|
def _cur(self):
|
|
"""Yield a cursor while holding the driver lock (thread-safe execute→fetch)."""
|
|
with self._lock:
|
|
cursor = self._connection.cursor()
|
|
try:
|
|
yield cursor
|
|
finally:
|
|
try:
|
|
cursor.close()
|
|
except Exception:
|
|
pass
|
|
|
|
# ── Schema introspection ──────────────────────────────────────────────────
|
|
|
|
def get_databases(self) -> list:
|
|
with self._cur() as c:
|
|
c.execute(
|
|
"SELECT datname FROM pg_database "
|
|
"WHERE datistemplate = false ORDER BY datname"
|
|
)
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_tables(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT t.table_name, t.table_schema,
|
|
COALESCE(s.n_live_tup, 0),
|
|
0,
|
|
'',
|
|
COALESCE(obj_description(
|
|
(quote_ident(t.table_schema)||'.'||quote_ident(t.table_name))::regclass,
|
|
'pg_class'), '')
|
|
FROM information_schema.tables t
|
|
LEFT JOIN pg_stat_user_tables s
|
|
ON s.schemaname = t.table_schema AND s.relname = t.table_name
|
|
WHERE t.table_schema NOT IN ('pg_catalog','information_schema')
|
|
AND t.table_type = 'BASE TABLE'
|
|
ORDER BY t.table_name
|
|
""")
|
|
return [TableInfo(name=r[0], schema=r[1], row_count=r[2],
|
|
size_bytes=r[3], engine=r[4], comment=r[5])
|
|
for r in c.fetchall()]
|
|
|
|
def get_views(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT table_name FROM information_schema.views
|
|
WHERE table_schema NOT IN ('pg_catalog','information_schema')
|
|
ORDER BY table_name
|
|
""")
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_columns(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT c.column_name, c.data_type, c.is_nullable,
|
|
c.column_default,
|
|
(SELECT true FROM information_schema.table_constraints tc
|
|
JOIN information_schema.key_column_usage kcu
|
|
ON tc.constraint_name = kcu.constraint_name
|
|
WHERE tc.constraint_type = 'PRIMARY KEY'
|
|
AND kcu.table_name = c.table_name
|
|
AND kcu.column_name = c.column_name
|
|
LIMIT 1) IS NOT NULL,
|
|
false
|
|
FROM information_schema.columns c
|
|
WHERE c.table_name = %s
|
|
ORDER BY c.ordinal_position
|
|
""", (table,))
|
|
return [ColumnInfo(name=r[0], data_type=r[1],
|
|
nullable=(r[2] == "YES"), default=r[3],
|
|
is_primary_key=bool(r[4]),
|
|
is_foreign_key=bool(r[5]))
|
|
for r in c.fetchall()]
|
|
|
|
def get_indexes(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT i.relname, ix.indisunique,
|
|
array_agg(a.attname ORDER BY k.n) AS cols
|
|
FROM pg_class t
|
|
JOIN pg_index ix ON t.oid = ix.indrelid
|
|
JOIN pg_class i ON i.oid = ix.indexrelid
|
|
JOIN unnest(ix.indkey) WITH ORDINALITY AS k(attnum, n)
|
|
ON TRUE
|
|
JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum
|
|
WHERE t.relname = %s
|
|
GROUP BY i.relname, ix.indisunique
|
|
ORDER BY i.relname
|
|
""", (table,))
|
|
return [IndexInfo(name=r[0], columns=list(r[2]),
|
|
is_unique=bool(r[1]))
|
|
for r in c.fetchall()]
|
|
|
|
def get_foreign_keys(self, database: str, table: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT tc.constraint_name, kcu.column_name,
|
|
ccu.table_name, ccu.column_name,
|
|
rc.update_rule, rc.delete_rule
|
|
FROM information_schema.table_constraints tc
|
|
JOIN information_schema.key_column_usage kcu
|
|
ON tc.constraint_name = kcu.constraint_name
|
|
JOIN information_schema.constraint_column_usage ccu
|
|
ON ccu.constraint_name = tc.constraint_name
|
|
JOIN information_schema.referential_constraints rc
|
|
ON rc.constraint_name = tc.constraint_name
|
|
WHERE tc.constraint_type = 'FOREIGN KEY'
|
|
AND tc.table_name = %s
|
|
""", (table,))
|
|
return [ForeignKeyInfo(name=r[0], column=r[1],
|
|
ref_table=r[2], ref_column=r[3],
|
|
on_update=r[4] or "", on_delete=r[5] or "")
|
|
for r in c.fetchall()]
|
|
|
|
def get_table_ddl(self, database: str, table: str) -> str:
|
|
cols = self.get_columns(database, table)
|
|
lines = [f"CREATE TABLE {table} ("]
|
|
col_defs = []
|
|
for c in cols:
|
|
d = f" {c.name} {c.data_type}"
|
|
if not c.nullable: d += " NOT NULL"
|
|
if c.default: d += f" DEFAULT {c.default}"
|
|
col_defs.append(d)
|
|
lines.append(",\n".join(col_defs))
|
|
lines.append(");")
|
|
return "\n".join(lines)
|
|
|
|
def get_functions(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT routine_name FROM information_schema.routines
|
|
WHERE routine_type = 'FUNCTION'
|
|
AND routine_schema NOT IN ('pg_catalog','information_schema')
|
|
ORDER BY routine_name
|
|
""")
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_stored_procedures(self, database: str) -> list:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT routine_name FROM information_schema.routines
|
|
WHERE routine_type = 'PROCEDURE'
|
|
AND routine_schema NOT IN ('pg_catalog','information_schema')
|
|
ORDER BY routine_name
|
|
""")
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
def get_triggers(self, database: str, table: str = "") -> list:
|
|
with self._cur() as c:
|
|
if table:
|
|
c.execute("""
|
|
SELECT trigger_name FROM information_schema.triggers
|
|
WHERE event_object_table = %s ORDER BY trigger_name
|
|
""", (table,))
|
|
else:
|
|
c.execute(
|
|
"SELECT trigger_name FROM information_schema.triggers "
|
|
"ORDER BY trigger_name"
|
|
)
|
|
return [r[0] for r in c.fetchall()]
|
|
|
|
# ── Query execution ───────────────────────────────────────────────────────
|
|
|
|
def execute_query(self, sql: str, params: Optional[tuple] = None) -> tuple:
|
|
with self._cur() as c:
|
|
c.execute(sql, params)
|
|
if c.description:
|
|
cols = [d[0] for d in c.description]
|
|
rows = c.fetchall()
|
|
return cols, rows, len(rows)
|
|
return [], [], c.rowcount
|
|
|
|
def execute_script(self, sql: str) -> list:
|
|
results = []
|
|
stmts = [s.strip() for s in sql.split(';') if s.strip()]
|
|
with self._cur() as c:
|
|
for stmt in stmts:
|
|
try:
|
|
c.execute(stmt)
|
|
if c.description:
|
|
cols = [d[0] for d in c.description]
|
|
rows = c.fetchall()
|
|
results.append((cols, rows, len(rows), ""))
|
|
else:
|
|
results.append(([], [], c.rowcount,
|
|
f"{c.rowcount} row(s) affected"))
|
|
except Exception as e:
|
|
results.append(([], [], 0, f"Error: {e}"))
|
|
return results
|
|
|
|
# ── Table data CRUD ───────────────────────────────────────────────────────
|
|
|
|
def get_table_data(self, database: str, table: str,
|
|
where: str = "", order_by: str = "",
|
|
limit: int = 1000, offset: int = 0) -> tuple:
|
|
sql = f'SELECT * FROM "{table}"'
|
|
if where: sql += f" WHERE {where}"
|
|
if order_by: sql += f" ORDER BY {order_by}"
|
|
sql += f" LIMIT {limit} OFFSET {offset}"
|
|
return self.execute_query(sql)
|
|
|
|
def get_table_row_count(self, database: str, table: str, where: str = "") -> int:
|
|
sql = f'SELECT COUNT(*) FROM "{table}"'
|
|
if where: sql += f" WHERE {where}"
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return c.fetchone()[0]
|
|
|
|
def insert_row(self, database: str, table: str, data: dict) -> bool:
|
|
cols = ", ".join(f'"{c}"' for c in data)
|
|
ph = ", ".join(["%s"] * len(data))
|
|
with self._cur() as c:
|
|
c.execute(f'INSERT INTO "{table}" ({cols}) VALUES ({ph})',
|
|
tuple(data.values()))
|
|
return True
|
|
|
|
@staticmethod
|
|
def _where(where: dict) -> tuple:
|
|
parts, params = [], []
|
|
for col, val in where.items():
|
|
if val is None:
|
|
parts.append(f'"{col}" IS NULL')
|
|
else:
|
|
parts.append(f'"{col}" = %s')
|
|
params.append(val)
|
|
return " AND ".join(parts), params
|
|
|
|
def update_row(self, database: str, table: str, data: dict, where: dict) -> bool:
|
|
set_cl = ", ".join(f'"{c}" = %s' for c in data)
|
|
where_cl, where_params = self._where(where)
|
|
with self._cur() as c:
|
|
c.execute(f'UPDATE "{table}" SET {set_cl} WHERE {where_cl}',
|
|
tuple(data.values()) + tuple(where_params))
|
|
return True
|
|
|
|
def delete_row(self, database: str, table: str, where: dict) -> bool:
|
|
where_cl, where_params = self._where(where)
|
|
with self._cur() as c:
|
|
c.execute(f'DELETE FROM "{table}" WHERE {where_cl}',
|
|
tuple(where_params))
|
|
return True
|
|
|
|
# ── Server tools ──────────────────────────────────────────────────────────
|
|
|
|
def explain_query(self, sql: str) -> tuple:
|
|
with self._cur() as c:
|
|
c.execute(f"EXPLAIN ANALYZE {sql}")
|
|
return ["Plan"], c.fetchall()
|
|
|
|
def get_process_list(self) -> tuple:
|
|
with self._cur() as c:
|
|
c.execute("""
|
|
SELECT pid, usename, application_name, client_addr,
|
|
state, query, query_start
|
|
FROM pg_stat_activity WHERE state IS NOT NULL
|
|
ORDER BY query_start DESC NULLS LAST
|
|
""")
|
|
cols = [d[0] for d in c.description]
|
|
return cols, c.fetchall()
|
|
|
|
def kill_process(self, process_id: int) -> bool:
|
|
with self._cur() as c:
|
|
c.execute("SELECT pg_terminate_backend(%s)", (process_id,))
|
|
return True
|
|
|
|
# ── Table designer ────────────────────────────────────────────────────────
|
|
|
|
def add_column(self, _database: str, table: str, col_name: str,
|
|
col_type: str, nullable: bool = True,
|
|
default=None) -> bool:
|
|
default_clause = f" DEFAULT {default}" if default is not None and default != "" else ""
|
|
null_clause = "" if nullable else " NOT NULL"
|
|
sql = (f'ALTER TABLE "{table}" '
|
|
f'ADD COLUMN "{col_name}" {col_type}{default_clause}{null_clause}')
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return True
|
|
|
|
def drop_column(self, _database: str, table: str, col_name: str) -> bool:
|
|
sql = f'ALTER TABLE "{table}" DROP COLUMN "{col_name}"'
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return True
|
|
|
|
def rename_column(self, _database: str, table: str,
|
|
old_name: str, new_name: str) -> bool:
|
|
sql = f'ALTER TABLE "{table}" RENAME COLUMN "{old_name}" TO "{new_name}"'
|
|
with self._cur() as c:
|
|
c.execute(sql)
|
|
return True
|