# ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ========= Copyright 2026 @ Strukto.AI All Rights Reserved. ========= import json import asyncpg async def list_schemas(conn: asyncpg.Connection, allowlist: list[str] | None) -> list[str]: rows = await conn.fetch( "SELECT schema_name FROM information_schema.schemata " "WHERE schema_name NOT IN ('pg_catalog', 'information_schema') " "AND schema_name NOT LIKE 'pg_%' " "ORDER BY schema_name") names = [r["schema_name"] for r in rows] if allowlist is not None: names = [n for n in names if n in allowlist] return names async def list_tables(conn: asyncpg.Connection, schema: str) -> list[str]: rows = await conn.fetch( "SELECT table_name FROM information_schema.tables " "WHERE table_schema = $1 AND table_type = 'BASE TABLE' " "ORDER BY table_name", schema) return [r["table_name"] for r in rows] async def list_views(conn: asyncpg.Connection, schema: str) -> list[str]: rows = await conn.fetch( "SELECT table_name FROM information_schema.views " "WHERE table_schema = $1 " "ORDER BY table_name", schema) return [r["table_name"] for r in rows] async def list_matviews(conn: asyncpg.Connection, schema: str) -> list[str]: rows = await conn.fetch( "SELECT matviewname AS name FROM pg_matviews " "WHERE schemaname = $1 " "ORDER BY matviewname", schema) return [r["name"] for r in rows] async def count_rows(conn: asyncpg.Connection, schema: str, name: str) -> int: return await conn.fetchval(f'SELECT COUNT(*) FROM "{schema}"."{name}"') async def estimate_size(conn: asyncpg.Connection, schema: str, name: str) -> tuple[int, int]: plan = await conn.fetchval( f'EXPLAIN (FORMAT JSON) SELECT * FROM "{schema}"."{name}"') if isinstance(plan, str): plan = json.loads(plan) top = plan[0]["Plan"] return int(top.get("Plan Rows", 0)), int(top.get("Plan Width", 0)) async def estimated_row_count(conn: asyncpg.Connection, schema: str, name: str) -> int: val = await conn.fetchval( "SELECT reltuples::bigint FROM pg_class c " "JOIN pg_namespace n ON c.relnamespace = n.oid " "WHERE n.nspname = $1 AND c.relname = $2", schema, name) return int(val) if val is not None else 0 async def table_size_bytes(conn: asyncpg.Connection, schema: str, name: str) -> int: val = await conn.fetchval( "SELECT pg_total_relation_size(c.oid) FROM pg_class c " "JOIN pg_namespace n ON c.relnamespace = n.oid " "WHERE n.nspname = $1 AND c.relname = $2", schema, name) return int(val) if val is not None else 0 async def fetch_rows(conn: asyncpg.Connection, schema: str, name: str, *, limit: int, offset: int) -> list[dict]: rows = await conn.fetch( f'SELECT * FROM "{schema}"."{name}" LIMIT $1 OFFSET $2', limit, offset) return [dict(r) for r in rows] async def fetch_columns(conn: asyncpg.Connection, schema: str, name: str) -> list[dict]: rows = await conn.fetch( "SELECT column_name, data_type, is_nullable " "FROM information_schema.columns " "WHERE table_schema = $1 AND table_name = $2 " "ORDER BY ordinal_position", schema, name) return [{ "name": r["column_name"], "type": r["data_type"], "nullable": r["is_nullable"] == "YES", } for r in rows] async def fetch_primary_key(conn: asyncpg.Connection, schema: str, name: str) -> list[str]: rows = await conn.fetch( "SELECT kcu.column_name " "FROM information_schema.table_constraints tc " "JOIN information_schema.key_column_usage kcu " " ON tc.constraint_name = kcu.constraint_name " " AND tc.table_schema = kcu.table_schema " "WHERE tc.constraint_type = 'PRIMARY KEY' " " AND tc.table_schema = $1 AND tc.table_name = $2 " "ORDER BY kcu.ordinal_position", schema, name) return [r["column_name"] for r in rows] async def fetch_foreign_keys(conn: asyncpg.Connection, schema: str, name: str) -> list[dict]: rows = await conn.fetch( "SELECT con.conname AS constraint_name, " " a.attname AS from_column, " " af.attname AS to_column, " " k.ord, " " nf.nspname AS to_schema, " " cf.relname AS to_table " "FROM pg_constraint con " "JOIN pg_class c ON c.oid = con.conrelid " "JOIN pg_namespace n ON n.oid = c.relnamespace " "JOIN pg_class cf ON cf.oid = con.confrelid " "JOIN pg_namespace nf ON nf.oid = cf.relnamespace " "JOIN unnest(con.conkey) WITH ORDINALITY AS k(attnum, ord) ON TRUE " "JOIN unnest(con.confkey) WITH ORDINALITY AS kf(attnum, ord) " " ON kf.ord = k.ord " "JOIN pg_attribute a " " ON a.attrelid = con.conrelid AND a.attnum = k.attnum " "JOIN pg_attribute af " " ON af.attrelid = con.confrelid AND af.attnum = kf.attnum " "WHERE con.contype = 'f' AND n.nspname = $1 AND c.relname = $2 " "ORDER BY con.conname, k.ord", schema, name) grouped: dict[str, dict] = {} for r in rows: cn = r["constraint_name"] if cn not in grouped: grouped[cn] = { "columns": [], "references": { "schema": r["to_schema"], "table": r["to_table"], "columns": [], }, } grouped[cn]["columns"].append(r["from_column"]) grouped[cn]["references"]["columns"].append(r["to_column"]) return list(grouped.values()) async def fetch_indexes(conn: asyncpg.Connection, schema: str, name: str) -> list[dict]: rows = await conn.fetch( "SELECT i.relname AS name, " " ix.indisunique AS unique, " " array_agg(a.attname ORDER BY x.ord) AS columns " "FROM pg_class t " "JOIN pg_namespace n ON t.relnamespace = n.oid " "JOIN pg_index ix ON ix.indrelid = t.oid " "JOIN pg_class i ON i.oid = ix.indexrelid " "JOIN unnest(ix.indkey) WITH ORDINALITY AS x(attnum, ord) ON TRUE " "JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = x.attnum " "WHERE n.nspname = $1 AND t.relname = $2 " "GROUP BY i.relname, ix.indisunique " "ORDER BY i.relname", schema, name) return [{ "name": r["name"], "columns": list(r["columns"]), "unique": r["unique"], } for r in rows] async def fetch_all_relationships(conn: asyncpg.Connection, schemas: list[str]) -> list[dict]: if not schemas: return [] rows = await conn.fetch( "SELECT con.conname AS constraint_name, " " n.nspname AS from_schema, " " c.relname AS from_table, " " a.attname AS from_column, " " af.attname AS to_column, " " k.ord, " " nf.nspname AS to_schema, " " cf.relname AS to_table " "FROM pg_constraint con " "JOIN pg_class c ON c.oid = con.conrelid " "JOIN pg_namespace n ON n.oid = c.relnamespace " "JOIN pg_class cf ON cf.oid = con.confrelid " "JOIN pg_namespace nf ON nf.oid = cf.relnamespace " "JOIN unnest(con.conkey) WITH ORDINALITY AS k(attnum, ord) ON TRUE " "JOIN unnest(con.confkey) WITH ORDINALITY AS kf(attnum, ord) " " ON kf.ord = k.ord " "JOIN pg_attribute a " " ON a.attrelid = con.conrelid AND a.attnum = k.attnum " "JOIN pg_attribute af " " ON af.attrelid = con.confrelid AND af.attnum = kf.attnum " "WHERE con.contype = 'f' AND n.nspname = ANY($1::text[]) " "ORDER BY n.nspname, c.relname, con.conname, k.ord", schemas) grouped: dict[tuple[str, str, str], dict] = {} for r in rows: key = (r["from_schema"], r["from_table"], r["constraint_name"]) if key not in grouped: grouped[key] = { "from": { "schema": r["from_schema"], "table": r["from_table"], "columns": [], }, "to": { "schema": r["to_schema"], "table": r["to_table"], "columns": [], }, "kind": "many_to_one", } grouped[key]["from"]["columns"].append(r["from_column"]) grouped[key]["to"]["columns"].append(r["to_column"]) return list(grouped.values())