pretensor 0.1.0__py3-none-any.whl

pretensor/connectors/registry.py

@@ -0,0 +1,93 @@
+"""Registry mapping connection types to connector implementations."""
+
+from __future__ import annotations
+
+from importlib import import_module
+
+from pretensor.connectors.base import BaseConnector
+from pretensor.introspection.models.config import ConnectionConfig, DatabaseType
+
+_POSTGRES_EXTRA_HINT = (
+    "Install the Postgres extra: pip install 'pretensor[postgres]' "
+    "(or pip install psycopg2-binary)."
+)
+
+_SNOWFLAKE_EXTRA_HINT = (
+    "Install the Snowflake extra: pip install 'pretensor[snowflake]' "
+    "(or pip install snowflake-sqlalchemy)."
+)
+
+_MYSQL_EXTRA_HINT = (
+    "Install the MySQL extra: pip install 'pretensor[mysql]' (or pip install PyMySQL)."
+)
+
+
+def _postgres_connector_class() -> type[BaseConnector]:
+    try:
+        # psycopg2 is an optional dependency under the ``postgres`` extra. Import
+        # it lazily here (mirroring the other connectors) so a bare install can
+        # import pretensor and only fails — with a clear hint — when a Postgres
+        # connection is actually attempted.
+        import_module("psycopg2")
+    except ImportError as exc:
+        msg = f"Postgres connector requires psycopg2. {_POSTGRES_EXTRA_HINT}"
+        raise ImportError(msg) from exc
+    from pretensor.connectors.postgres import PostgresConnector
+
+    return PostgresConnector
+
+
+def _snowflake_connector_class() -> type[BaseConnector]:
+    try:
+        # Optional dep: string path keeps pyright happy when the extra is not installed.
+        import_module("snowflake.sqlalchemy")
+    except ImportError as exc:
+        msg = f"Snowflake connector requires snowflake-sqlalchemy. {_SNOWFLAKE_EXTRA_HINT}"
+        raise ImportError(msg) from exc
+    from pretensor.connectors.snowflake import SnowflakeConnector
+
+    return SnowflakeConnector
+
+
+def _bigquery_connector_class() -> type[BaseConnector]:
+    try:
+        import_module("google.cloud.bigquery")
+    except ImportError as exc:
+        from pretensor.connectors.bigquery import _BIGQUERY_EXTRA_HINT
+
+        msg = (
+            f"BigQuery connector requires google-cloud-bigquery. {_BIGQUERY_EXTRA_HINT}"
+        )
+        raise ImportError(msg) from exc
+    from pretensor.connectors.bigquery import BigQueryConnector
+
+    return BigQueryConnector
+
+
+def _mysql_connector_class() -> type[BaseConnector]:
+    try:
+        import_module("pymysql")
+    except ImportError as exc:
+        msg = f"MySQL connector requires PyMySQL. {_MYSQL_EXTRA_HINT}"
+        raise ImportError(msg) from exc
+    from pretensor.connectors.mysql import MySQLConnector
+
+    return MySQLConnector
+
+
+def get_connector(config: ConnectionConfig) -> BaseConnector:
+    """Return a connector instance for the given configuration."""
+    if config.type == DatabaseType.POSTGRES:
+        cls = _postgres_connector_class()
+        return cls(config)
+    if config.type == DatabaseType.SNOWFLAKE:
+        cls = _snowflake_connector_class()
+        return cls(config)
+    if config.type == DatabaseType.BIGQUERY:
+        cls = _bigquery_connector_class()
+        return cls(config)
+    if config.type == DatabaseType.MYSQL:
+        cls = _mysql_connector_class()
+        return cls(config)
+    msg = f"No connector registered for database type: {config.type}"
+    raise ValueError(msg)

pretensor/connectors/snapshot.py

@@ -0,0 +1,244 @@
+"""Snapshot diff — compare two schema snapshots for drift detection."""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+from pretensor.connectors.models import Column, SchemaSnapshot, Table, ViewDependency
+from pretensor.introspection.models.base import PretensorModel
+
+__all__ = [
+    "ChangeTarget",
+    "ChangeType",
+    "SchemaChange",
+    "diff_snapshots",
+]
+
+
+class ChangeType(StrEnum):
+    ADDED = "added"
+    REMOVED = "removed"
+    MODIFIED = "modified"
+
+
+class ChangeTarget(StrEnum):
+    TABLE = "table"
+    COLUMN = "column"
+    LINEAGE = "lineage"
+
+
+class SchemaChange(PretensorModel):
+    """A single difference between two schema snapshots."""
+
+    change_type: ChangeType
+    target: ChangeTarget
+    table_name: str
+    schema_name: str
+    column_name: str | None = None
+    details: str = ""
+
+
+def diff_snapshots(old: SchemaSnapshot, new: SchemaSnapshot) -> list[SchemaChange]:
+    """Compare two snapshots and return a list of changes.
+
+    Detects added/removed tables and added/removed/modified columns.
+    """
+    changes: list[SchemaChange] = []
+
+    old_tables = {(t.schema_name, t.name): t for t in old.tables}
+    new_tables = {(t.schema_name, t.name): t for t in new.tables}
+
+    old_keys = set(old_tables.keys())
+    new_keys = set(new_tables.keys())
+
+    for schema_name, table_name in sorted(new_keys - old_keys):
+        changes.append(
+            SchemaChange(
+                change_type=ChangeType.ADDED,
+                target=ChangeTarget.TABLE,
+                table_name=table_name,
+                schema_name=schema_name,
+                details=f"New table with {len(new_tables[(schema_name, table_name)].columns)} columns",
+            )
+        )
+
+    for schema_name, table_name in sorted(old_keys - new_keys):
+        changes.append(
+            SchemaChange(
+                change_type=ChangeType.REMOVED,
+                target=ChangeTarget.TABLE,
+                table_name=table_name,
+                schema_name=schema_name,
+            )
+        )
+
+    for key in sorted(old_keys & new_keys):
+        schema_name, table_name = key
+        old_table = old_tables[key]
+        new_table = new_tables[key]
+
+        old_cols = {c.name: c for c in old_table.columns}
+        new_cols = {c.name: c for c in new_table.columns}
+
+        for col_name in sorted(set(new_cols) - set(old_cols)):
+            col = new_cols[col_name]
+            changes.append(
+                SchemaChange(
+                    change_type=ChangeType.ADDED,
+                    target=ChangeTarget.COLUMN,
+                    table_name=table_name,
+                    schema_name=schema_name,
+                    column_name=col_name,
+                    details=f"type={col.data_type}, nullable={col.nullable}",
+                )
+            )
+
+        for col_name in sorted(set(old_cols) - set(new_cols)):
+            changes.append(
+                SchemaChange(
+                    change_type=ChangeType.REMOVED,
+                    target=ChangeTarget.COLUMN,
+                    table_name=table_name,
+                    schema_name=schema_name,
+                    column_name=col_name,
+                )
+            )
+
+        for col_name in sorted(set(old_cols) & set(new_cols)):
+            old_col = old_cols[col_name]
+            new_col = new_cols[col_name]
+            diffs = _column_diffs(old_col, new_col)
+            if diffs:
+                changes.append(
+                    SchemaChange(
+                        change_type=ChangeType.MODIFIED,
+                        target=ChangeTarget.COLUMN,
+                        table_name=table_name,
+                        schema_name=schema_name,
+                        column_name=col_name,
+                        details=", ".join(diffs),
+                    )
+                )
+
+        table_diffs = _table_diffs(old_table, new_table)
+        if table_diffs:
+            changes.append(
+                SchemaChange(
+                    change_type=ChangeType.MODIFIED,
+                    target=ChangeTarget.TABLE,
+                    table_name=table_name,
+                    schema_name=schema_name,
+                    details=", ".join(table_diffs),
+                )
+            )
+
+    changes.extend(_lineage_diffs(old.view_dependencies, new.view_dependencies))
+
+    return changes
+
+
+def _lineage_key(dep: ViewDependency) -> tuple[str, str, str, str, str]:
+    return (
+        dep.source_schema,
+        dep.source_table,
+        dep.target_schema,
+        dep.target_table,
+        dep.lineage_type,
+    )
+
+
+def _lineage_diffs(
+    old_deps: list[ViewDependency], new_deps: list[ViewDependency]
+) -> list[SchemaChange]:
+    """Detect added and removed lineage edges between two snapshots."""
+    old_map = {_lineage_key(d): d for d in old_deps}
+    new_map = {_lineage_key(d): d for d in new_deps}
+    changes: list[SchemaChange] = []
+    for key in sorted(set(new_map) - set(old_map)):
+        d = new_map[key]
+        changes.append(
+            SchemaChange(
+                change_type=ChangeType.ADDED,
+                target=ChangeTarget.LINEAGE,
+                schema_name=d.source_schema,
+                table_name=d.source_table,
+                details=(
+                    f"{d.lineage_type}: {d.source_schema}.{d.source_table}"
+                    f" → {d.target_schema}.{d.target_table}"
+                ),
+            )
+        )
+    for key in sorted(set(old_map) - set(new_map)):
+        d = old_map[key]
+        changes.append(
+            SchemaChange(
+                change_type=ChangeType.REMOVED,
+                target=ChangeTarget.LINEAGE,
+                schema_name=d.source_schema,
+                table_name=d.source_table,
+                details=(
+                    f"{d.lineage_type}: {d.source_schema}.{d.source_table}"
+                    f" → {d.target_schema}.{d.target_table}"
+                ),
+            )
+        )
+    return changes
+
+
+def _column_diffs(old: Column, new: Column) -> list[str]:
+    """Compare two Column objects on structural and catalog-stat fields.
+
+    ``parent_column`` and ``is_array`` are intentionally excluded: they are
+    set once during initial introspection (e.g. BigQuery nested fields) and
+    are not expected to change across re-indexes of the same table.
+    """
+    structural_fields = (
+        "data_type",
+        "nullable",
+        "is_primary_key",
+        "is_foreign_key",
+        "default_value",
+        "comment",
+        "is_indexed",
+        "check_constraints",
+        "ordinal_position",
+        "most_common_values",
+        "histogram_bounds",
+        "stats_correlation",
+    )
+    diffs: list[str] = []
+    for field in structural_fields:
+        old_val = getattr(old, field)
+        new_val = getattr(new, field)
+        if old_val != new_val:
+            diffs.append(f"{field}: {old_val!r} → {new_val!r}")
+    return diffs
+
+
+def _table_diffs(old: Table, new: Table) -> list[str]:
+    """Compare table-level metadata (not column sets)."""
+    diffs: list[str] = []
+    for field in (
+        "comment",
+        "table_type",
+        "row_count",
+        "seq_scan_count",
+        "idx_scan_count",
+        "insert_count",
+        "update_count",
+        "delete_count",
+        "is_partitioned",
+        "partition_key",
+        "grants",
+        "access_read_count",
+        "access_write_count",
+        "days_since_last_access",
+        "potentially_unused",
+        "table_bytes",
+        "clustering_key",
+    ):
+        old_val = getattr(old, field)
+        new_val = getattr(new, field)
+        if old_val != new_val:
+            diffs.append(f"{field}: {old_val!r} → {new_val!r}")
+    return diffs