"""PostgreSQL engine — optional dependency: psycopg[binary]>=3.1.0.""" import time from mcp_database_universal.engines.base import ( BaseEngine, DBInfo, ColumnInfo, TableInfo, TableDetail, IndexInfo, ForeignKeyInfo, TableStats, QueryResult, _named_to_pyformat, ) class PostgresEngine(BaseEngine): def __init__( self, host: str = "localhost", port: int = 5432, database: str = "postgres", user: str = "", password: str = "", read_only: bool = True, ): self.host = host self.port = port self.database = database self.user = user self.password = password self._read_only = read_only self._conn = None async def connect(self) -> None: try: import psycopg except ImportError: raise RuntimeError( "PostgreSQL engine requires psycopg. " "Install with: pip install 'mcp-database-universal[postgres]'" ) self._conn = await psycopg.AsyncConnection.connect( host=self.host, port=self.port, dbname=self.database, user=self.user, password=self.password, autocommit=True, ) async def disconnect(self) -> None: if self._conn: await self._conn.close() self._conn = None def _ensure_conn(self): if self._conn is None: raise RuntimeError("Not connected. Call connect() first.") return self._conn async def get_db_info(self) -> DBInfo: conn = self._ensure_conn() cur = await conn.execute("SELECT version()") row = await cur.fetchone() version = row[0] if row else "unknown" try: cur = await conn.execute( "SELECT pg_size_pretty(pg_database_size(current_database()))" ) row2 = await cur.fetchone() size = row2[0] if row2 else "unknown" except Exception: size = "unknown" return DBInfo(engine="postgresql", version=version, name=self.database, size_approx=size) async def get_tables(self) -> list[TableInfo]: conn = self._ensure_conn() cur = await conn.execute(""" SELECT t.table_name FROM information_schema.tables t WHERE t.table_schema = 'public' AND t.table_type = 'BASE TABLE' ORDER BY t.table_name """) rows = await cur.fetchall() tables = [] for row in rows: name = row[0] try: cur = await conn.execute(f'SELECT COUNT(*) FROM "{name}"') cnt = await cur.fetchone() row_count = cnt[0] if cnt else 0 except Exception: row_count = 0 try: cur = await conn.execute(""" SELECT COUNT(*) FROM information_schema.columns WHERE table_schema = 'public' AND table_name = %s """, [name]) cols = await cur.fetchone() col_count = cols[0] if cols else 0 except Exception: col_count = 0 tables.append(TableInfo( name=name, row_count=row_count, column_count=col_count, )) return tables async def get_table_detail(self, table: str) -> TableDetail: conn = self._ensure_conn() cur = await conn.execute(""" SELECT c.column_name, c.data_type, c.is_nullable, c.column_default, CASE WHEN pk.column_name IS NOT NULL THEN true ELSE false END as is_pk FROM information_schema.columns c LEFT JOIN ( SELECT ku.column_name FROM information_schema.table_constraints tc JOIN information_schema.key_column_usage ku ON tc.constraint_name = ku.constraint_name WHERE tc.table_name = %s AND tc.constraint_type = 'PRIMARY KEY' ) pk ON c.column_name = pk.column_name WHERE c.table_schema = 'public' AND c.table_name = %s ORDER BY c.ordinal_position """, [table, table]) col_rows = await cur.fetchall() columns = [] pk_name = None for c in col_rows: is_pk = c[4] if is_pk: pk_name = c[0] columns.append(ColumnInfo( name=c[0], type=c[1], nullable=(c[2] == "YES"), default=c[3], is_primary_key=is_pk, )) cur = await conn.execute(""" SELECT kcu.column_name, ccu.table_name AS foreign_table_name, ccu.column_name AS foreign_column_name 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 tc.constraint_name = ccu.constraint_name WHERE tc.constraint_type = 'FOREIGN KEY' AND tc.table_name = %s """, [table]) fk_rows = await cur.fetchall() foreign_keys = [] fk_cols = set() for fk in fk_rows: foreign_keys.append(ForeignKeyInfo( column=fk[0], references_table=fk[1], references_column=fk[2], )) fk_cols.add(fk[0]) for col in columns: if col.name in fk_cols: col.is_foreign_key = True for fk in foreign_keys: if fk.column == col.name: col.foreign_key_table = fk.references_table col.foreign_key_column = fk.references_column break cur = await conn.execute(""" SELECT indexname, indexdef FROM pg_indexes WHERE tablename = %s AND schemaname = 'public' """, [table]) idx_rows = await cur.fetchall() indexes = [] for idx in idx_rows: name = idx[0] unique = "UNIQUE" in (idx[1] or "").upper() indexes.append(IndexInfo(name=name, columns=[], unique=unique)) try: cur = await conn.execute(f'SELECT * FROM "{table}" LIMIT 5') sample_rows = await cur.fetchall() if sample_rows: cols_d = [d.name for d in cur.description] if cur.description else [] sample_data = [dict(zip(cols_d, row)) for row in sample_rows] else: sample_data = [] except Exception: sample_data = [] stats = await self.get_table_stats(table) return TableDetail( name=table, columns=columns, indexes=indexes, primary_key=pk_name, foreign_keys=foreign_keys, sample_data=sample_data, stats=stats, ) async def execute_query(self, sql: str, params: dict | None = None) -> QueryResult: conn = self._ensure_conn() start = time.monotonic() try: sql, params = self._translate_params(sql, params) if params: cur = await conn.execute(sql, params) else: cur = await conn.execute(sql) rows = await cur.fetchall() columns = [desc.name for desc in cur.description] if cur.description else [] elapsed = int((time.monotonic() - start) * 1000) result_rows = [dict(zip(columns, row)) for row in rows] return QueryResult( columns=columns, rows=result_rows, row_count=len(result_rows), truncated=False, execution_time_ms=elapsed, sql=sql, ) except Exception as e: elapsed = int((time.monotonic() - start) * 1000) return QueryResult( sql=sql, execution_time_ms=elapsed, warning=f"Error: {str(e)}", ) async def get_sample_data(self, table: str, limit: int = 5) -> QueryResult: return await self.execute_query(f'SELECT * FROM "{table}" LIMIT {limit}') async def get_table_stats(self, table: str) -> TableStats: conn = self._ensure_conn() try: cur = await conn.execute(f'SELECT COUNT(*) FROM "{table}"') row = await cur.fetchone() row_count = row[0] if row else 0 except Exception: row_count = 0 try: cur = await conn.execute( "SELECT pg_size_pretty(pg_total_relation_size(%s))", [table] ) row = await cur.fetchone() total_size = row[0] if row else "unknown" except Exception: total_size = "unknown" return TableStats( row_count=row_count, total_size=total_size, ) def is_read_only(self) -> bool: return self._read_only