ThothII/harness/tht/db/introspect.py

from datetime import UTC, datetime

from sqlalchemy import Engine, text

from tht.mschema.models import (
    ColumnPhysical,
    ForeignKey,
    Index,
    PhysicalSchema,
    TablePhysical,
)

# Query adattate da thoth_sqldb2 (Apache 2.0) — vedi src/tht/vendor/VENDORED.md.

_TABLES_Q = text("""
    SELECT c.relname AS table_name,
           COALESCE(d.description, '') AS comment,
           GREATEST(c.reltuples::bigint, 0) AS row_count
    FROM pg_class c
    JOIN pg_namespace n ON n.oid = c.relnamespace
    LEFT JOIN pg_description d ON d.objoid = c.oid AND d.objsubid = 0
    WHERE c.relkind IN ('r', 'p') AND n.nspname = :schema
    ORDER BY c.relname
""")

_COLUMNS_Q = text("""
    SELECT a.attname AS column_name,
           format_type(a.atttypid, a.atttypmod) AS data_type,
           (NOT a.attnotnull) AS is_nullable,
           pg_get_expr(d.adbin, d.adrelid) AS column_default,
           COALESCE(pgd.description, '') AS comment,
           (ty.typtype = 'e') AS is_enum,
           EXISTS (
               SELECT 1 FROM pg_index i
               WHERE i.indrelid = c.oid AND i.indisprimary AND a.attnum = ANY (i.indkey)
           ) AS is_pk
    FROM pg_class c
    JOIN pg_namespace n ON n.oid = c.relnamespace
    JOIN pg_attribute a ON a.attrelid = c.oid
    JOIN pg_type ty ON ty.oid = a.atttypid
    LEFT JOIN pg_attrdef d ON d.adrelid = c.oid AND d.adnum = a.attnum
    LEFT JOIN pg_description pgd ON pgd.objoid = c.oid AND pgd.objsubid = a.attnum
    WHERE c.relname = :table_name AND n.nspname = :schema
      AND a.attnum > 0 AND NOT a.attisdropped
    ORDER BY a.attnum
""")

_FOREIGN_KEYS_Q = text("""
    SELECT con.conname AS constraint_name,
           rel.relname AS source_table,
           a.attname AS source_column,
           frel.relname AS target_table,
           fa.attname AS target_column,
           src.ord
    FROM pg_constraint con
    JOIN pg_class rel ON rel.oid = con.conrelid
    JOIN pg_namespace ns ON ns.oid = rel.relnamespace
    JOIN pg_class frel ON frel.oid = con.confrelid
    JOIN unnest(con.conkey) WITH ORDINALITY AS src(attnum, ord) ON true
    JOIN pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = src.attnum
    JOIN unnest(con.confkey) WITH ORDINALITY AS dst(attnum, ord) ON dst.ord = src.ord
    JOIN pg_attribute fa ON fa.attrelid = con.confrelid AND fa.attnum = dst.attnum
    WHERE con.contype = 'f' AND ns.nspname = :schema
    ORDER BY rel.relname, con.conname, src.ord
""")

_INDEXES_Q = text("""
    SELECT i.relname AS index_name,
           t.relname AS table_name,
           ix.indisunique AS is_unique,
           ix.indisprimary AS is_primary,
           am.amname AS index_type,
           array_agg(a.attname ORDER BY a.attnum) AS columns
    FROM pg_index ix
    JOIN pg_class i ON i.oid = ix.indexrelid
    JOIN pg_class t ON t.oid = ix.indrelid
    JOIN pg_namespace n ON n.oid = t.relnamespace
    JOIN pg_am am ON am.oid = i.relam
    JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY (ix.indkey)
    WHERE n.nspname = :schema
    GROUP BY i.relname, t.relname, ix.indisunique, ix.indisprimary, am.amname
    ORDER BY t.relname, i.relname
""")


class IntrospectionError(Exception):
    pass


def introspect(engine: Engine, database: str, schema: str) -> PhysicalSchema:
    with engine.connect() as conn:
        exists = conn.execute(
            text("SELECT 1 FROM pg_namespace WHERE nspname = :schema"), {"schema": schema}
        ).scalar()
        if not exists:
            raise IntrospectionError(f"Schema inesistente: {schema}")

        tables: dict[str, TablePhysical] = {}
        for trow in conn.execute(_TABLES_Q, {"schema": schema}):
            columns: dict[str, ColumnPhysical] = {}
            for crow in conn.execute(
                _COLUMNS_Q, {"table_name": trow.table_name, "schema": schema}
            ):
                columns[crow.column_name] = ColumnPhysical(
                    type=crow.data_type,
                    nullable=bool(crow.is_nullable),
                    pk=bool(crow.is_pk),
                    default=crow.column_default,
                    comment=crow.comment,
                    is_enum=bool(crow.is_enum),
                )
            tables[trow.table_name] = TablePhysical(
                comment=trow.comment, row_count=trow.row_count, columns=columns
            )

        # FK raggruppate per (tabella, constraint), ordinate per posizione
        grouped: dict[tuple[str, str], ForeignKey] = {}
        for row in conn.execute(_FOREIGN_KEYS_Q, {"schema": schema}):
            key = (row.source_table, row.constraint_name)
            fk = grouped.setdefault(
                key,
                ForeignKey(
                    columns=[], ref_table=row.target_table, ref_columns=[],
                    name=row.constraint_name,
                ),
            )
            fk.columns.append(row.source_column)
            fk.ref_columns.append(row.target_column)
        for (table_name, _), fk in grouped.items():
            if table_name in tables:
                tables[table_name].foreign_keys.append(fk)

        for row in conn.execute(_INDEXES_Q, {"schema": schema}):
            if row.table_name in tables:
                tables[row.table_name].indexes.append(
                    Index(
                        name=row.index_name,
                        columns=list(row.columns),
                        unique=bool(row.is_unique),
                        primary=bool(row.is_primary),
                        type=row.index_type,
                    )
                )

    return PhysicalSchema(
        database=database,
        schema=schema,
        introspected_at=datetime.now(UTC),
        tables=tables,
    )


def introspect_rest(client, database: str, schema: str) -> PhysicalSchema:
    """Introspezione via REST (rpc `list_tables`/`table_columns`/`table_comments`/
    `table_foreign_keys`). Limiti rispetto al diretto: niente indici (nessun rpc) e
    `is_enum` non disponibile (default False)."""
    tables: dict[str, TablePhysical] = {}
    for trow in client.list_tables(schema):
        if trow.get("type") != "TABLE":
            continue  # le viste sono fuori scope (come l'introspezione diretta)
        table_name = trow["table"]

        col_comments = {
            c["name"]: (c.get("comment") or "")
            for c in client.table_comments(schema, table_name)
            if c.get("object") == "COLUMN"
        }
        columns: dict[str, ColumnPhysical] = {}
        for crow in client.table_columns(schema, table_name):
            name = crow["column"]
            columns[name] = ColumnPhysical(
                type=crow["type"],
                nullable=bool(crow.get("nullable", True)),
                pk=bool(crow.get("pk", False)),
                default=crow.get("default"),
                comment=col_comments.get(name, ""),
            )

        foreign_keys: list[ForeignKey] = []
        for fk in client.table_foreign_keys(schema, table_name):
            foreign_keys.append(
                ForeignKey(
                    columns=fk.get("columns") or [fk["column"]],
                    ref_table=fk.get("ref_table") or fk["target_table"],
                    ref_columns=fk.get("ref_columns") or [fk["target_column"]],
                    name=fk.get("name", ""),
                )
            )

        tables[table_name] = TablePhysical(
            comment=trow.get("comment") or "",
            row_count=int(trow.get("rows") or 0),
            columns=columns,
            foreign_keys=foreign_keys,
        )

    return PhysicalSchema(
        database=database,
        schema=schema,
        introspected_at=datetime.now(UTC),
        tables=tables,
    )