Pārlūkot izejas kodu

feat: add postgresql enterprise connector

马小龙 2 dienas atpakaļ
vecāks
revīzija
57a481e3c7

+ 3 - 0
app/core/connectors/builtin/__init__.py

@@ -1,6 +1,7 @@
 """Built-in connectors registered exclusively through the public SDK."""
 
 from app.core.connectors.builtin.oracle import OracleConnector
+from app.core.connectors.builtin.postgresql import PostgreSQLConnector
 from app.core.connectors.builtin.rest_catalog import (
     RestCatalogConnector,
     SafeRestTransport,
@@ -10,6 +11,7 @@ from app.core.connectors.builtin.sqlserver import SqlServerConnector
 
 def register_builtin_connectors(registry, **dependencies):
     registry.register(OracleConnector(dependencies.get("connection_provider")))
+    registry.register(PostgreSQLConnector(dependencies.get("connection_provider")))
     registry.register(SqlServerConnector(dependencies.get("connection_provider")))
     transport = dependencies.get("rest_transport")
     if transport is None:
@@ -23,6 +25,7 @@ def register_builtin_connectors(registry, **dependencies):
 __all__ = [
     "register_builtin_connectors",
     "OracleConnector",
+    "PostgreSQLConnector",
     "SqlServerConnector",
     "RestCatalogConnector",
     "SafeRestTransport",

+ 68 - 0
app/core/connectors/builtin/postgresql.py

@@ -0,0 +1,68 @@
+"""PostgreSQL read-only catalog connector (driver and connection are optional)."""
+
+import importlib.util
+
+from app.core.connectors.builtin.database import (
+    DATABASE_CONFIG_SCHEMA,
+    ReadOnlyCatalogConnector,
+)
+from app.core.connectors.sdk import (
+    SDK_VERSION,
+    CompatibilityResult,
+    ConnectorManifest,
+    HealthResult,
+    validate_config,
+)
+
+POSTGRESQL_CATALOG_SQL = """
+SELECT c.table_schema AS schema_name, c.table_name AS asset_name,
+       CASE WHEN t.table_type = 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS asset_type,
+       c.column_name, c.ordinal_position, c.data_type, c.is_nullable,
+       c.column_default, d.description AS column_comment
+FROM information_schema.columns c
+JOIN information_schema.tables t
+  ON t.table_schema = c.table_schema AND t.table_name = c.table_name
+LEFT JOIN pg_catalog.pg_namespace n ON n.nspname = c.table_schema
+LEFT JOIN pg_catalog.pg_class r
+  ON r.relnamespace = n.oid AND r.relname = c.table_name
+LEFT JOIN pg_catalog.pg_description d
+  ON d.objoid = r.oid AND d.objsubid = c.ordinal_position
+WHERE t.table_type IN ('BASE TABLE', 'VIEW')
+  AND c.table_schema NOT IN ('pg_catalog', 'information_schema')
+ORDER BY c.table_schema, c.table_name, c.ordinal_position
+"""
+
+
+class PostgreSQLConnector(ReadOnlyCatalogConnector):
+    catalog_sql = POSTGRESQL_CATALOG_SQL
+    manifest = ConnectorManifest(
+        connector_id="postgresql",
+        version="1.0.0",
+        sdk_version=SDK_VERSION,
+        display_name="PostgreSQL",
+        capabilities=(
+            "discover",
+            "snapshot",
+            "incremental",
+            "cancel",
+            "resume",
+            "evidence",
+        ),
+        config_schema=DATABASE_CONFIG_SCHEMA,
+    )
+
+    def health(self, config):
+        validate_config(self.manifest.config_schema, config)
+        available = importlib.util.find_spec("psycopg2") is not None
+        return HealthResult(
+            "available" if available else "degraded",
+            "" if available else "optional driver unavailable",
+        )
+
+    def compatibility(self):
+        available = importlib.util.find_spec("psycopg2") is not None
+        return CompatibilityResult(
+            available,
+            self.manifest.version,
+            detail="" if available else "optional driver unavailable",
+        )

+ 3 - 0
deployment/app/core/connectors/builtin/__init__.py

@@ -1,6 +1,7 @@
 """Built-in connectors registered exclusively through the public SDK."""
 
 from app.core.connectors.builtin.oracle import OracleConnector
+from app.core.connectors.builtin.postgresql import PostgreSQLConnector
 from app.core.connectors.builtin.rest_catalog import (
     RestCatalogConnector,
     SafeRestTransport,
@@ -10,6 +11,7 @@ from app.core.connectors.builtin.sqlserver import SqlServerConnector
 
 def register_builtin_connectors(registry, **dependencies):
     registry.register(OracleConnector(dependencies.get("connection_provider")))
+    registry.register(PostgreSQLConnector(dependencies.get("connection_provider")))
     registry.register(SqlServerConnector(dependencies.get("connection_provider")))
     transport = dependencies.get("rest_transport")
     if transport is None:
@@ -23,6 +25,7 @@ def register_builtin_connectors(registry, **dependencies):
 __all__ = [
     "register_builtin_connectors",
     "OracleConnector",
+    "PostgreSQLConnector",
     "SqlServerConnector",
     "RestCatalogConnector",
     "SafeRestTransport",

+ 68 - 0
deployment/app/core/connectors/builtin/postgresql.py

@@ -0,0 +1,68 @@
+"""PostgreSQL read-only catalog connector (driver and connection are optional)."""
+
+import importlib.util
+
+from app.core.connectors.builtin.database import (
+    DATABASE_CONFIG_SCHEMA,
+    ReadOnlyCatalogConnector,
+)
+from app.core.connectors.sdk import (
+    SDK_VERSION,
+    CompatibilityResult,
+    ConnectorManifest,
+    HealthResult,
+    validate_config,
+)
+
+POSTGRESQL_CATALOG_SQL = """
+SELECT c.table_schema AS schema_name, c.table_name AS asset_name,
+       CASE WHEN t.table_type = 'VIEW' THEN 'VIEW' ELSE 'TABLE' END AS asset_type,
+       c.column_name, c.ordinal_position, c.data_type, c.is_nullable,
+       c.column_default, d.description AS column_comment
+FROM information_schema.columns c
+JOIN information_schema.tables t
+  ON t.table_schema = c.table_schema AND t.table_name = c.table_name
+LEFT JOIN pg_catalog.pg_namespace n ON n.nspname = c.table_schema
+LEFT JOIN pg_catalog.pg_class r
+  ON r.relnamespace = n.oid AND r.relname = c.table_name
+LEFT JOIN pg_catalog.pg_description d
+  ON d.objoid = r.oid AND d.objsubid = c.ordinal_position
+WHERE t.table_type IN ('BASE TABLE', 'VIEW')
+  AND c.table_schema NOT IN ('pg_catalog', 'information_schema')
+ORDER BY c.table_schema, c.table_name, c.ordinal_position
+"""
+
+
+class PostgreSQLConnector(ReadOnlyCatalogConnector):
+    catalog_sql = POSTGRESQL_CATALOG_SQL
+    manifest = ConnectorManifest(
+        connector_id="postgresql",
+        version="1.0.0",
+        sdk_version=SDK_VERSION,
+        display_name="PostgreSQL",
+        capabilities=(
+            "discover",
+            "snapshot",
+            "incremental",
+            "cancel",
+            "resume",
+            "evidence",
+        ),
+        config_schema=DATABASE_CONFIG_SCHEMA,
+    )
+
+    def health(self, config):
+        validate_config(self.manifest.config_schema, config)
+        available = importlib.util.find_spec("psycopg2") is not None
+        return HealthResult(
+            "available" if available else "degraded",
+            "" if available else "optional driver unavailable",
+        )
+
+    def compatibility(self):
+        available = importlib.util.find_spec("psycopg2") is not None
+        return CompatibilityResult(
+            available,
+            self.manifest.version,
+            detail="" if available else "optional driver unavailable",
+        )

+ 160 - 0
tests/test_postgresql_enterprise_connector.py

@@ -0,0 +1,160 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+import json
+import re
+
+import pytest
+
+from app.core.connectors.builtin import register_builtin_connectors
+from app.core.connectors.builtin.postgresql import PostgreSQLConnector
+from app.core.connectors.errors import ConnectorCancelledError
+from app.core.connectors.registry import ConnectorRegistry
+from app.core.connectors.sdk import OperationRequest
+
+
+_connections = []
+
+
+class _Rows:
+    def __init__(self, rows):
+        self.rows = rows
+
+    def mappings(self):
+        return self
+
+    def all(self):
+        return self.rows
+
+
+class _Connection:
+    def __init__(self, rows):
+        self.rows = rows
+        self.statement = None
+        self.parameters = None
+
+    def execute(self, statement, parameters):
+        self.statement = str(statement)
+        self.parameters = parameters
+        return _Rows(self.rows)
+
+
+@contextmanager
+def _provider(source_uid, purpose):
+    assert source_uid == "source-postgres"
+    assert purpose == "metadata_collection"
+    connection = _Connection(
+        [
+            {
+                "schema_name": "public",
+                "asset_name": "orders",
+                "asset_type": "TABLE",
+                "column_name": "id",
+                "ordinal_position": 1,
+                "data_type": "bigint",
+                "is_nullable": "NO",
+                "column_default": None,
+                "column_comment": "order key",
+            },
+            {
+                "schema_name": "internal",
+                "asset_name": "jobs",
+                "asset_type": "TABLE",
+                "column_name": "id",
+                "ordinal_position": 1,
+                "data_type": "bigint",
+                "is_nullable": "NO",
+                "column_default": None,
+                "column_comment": "job key",
+            },
+        ]
+    )
+    _connections.append(connection)
+    yield connection
+
+
+def _request(operation="discover", **values):
+    defaults = {
+        "source_uid": "source-postgres",
+        "operation": operation,
+        "config": {"credential_ref": "vault:postgresql/catalog"},
+        "scope": {"include_schemas": ["public"]},
+    }
+    defaults.update(values)
+    return OperationRequest(**defaults)
+
+
+def test_postgresql_registry_connector_collects_filtered_read_only_catalog():
+    _connections.clear()
+    registry = ConnectorRegistry()
+    register_builtin_connectors(registry, connection_provider=_provider)
+
+    connector = registry.resolve("postgresql", "1.0.0")
+    result = connector.discover(_request())
+
+    assert isinstance(connector, PostgreSQLConnector)
+    assert connector.manifest.display_name == "PostgreSQL"
+    assert connector.manifest.capabilities == (
+        "discover",
+        "snapshot",
+        "incremental",
+        "cancel",
+        "resume",
+        "evidence",
+    )
+    assert [record["asset_key"] for record in result.records] == [
+        "source-postgres:public.orders"
+    ]
+    assert result.checkpoint["snapshot_summary"]["record_count"] == 1
+    assert result.evidence["query_kind"] == "read_only_metadata"
+    assert "postgresql/catalog" not in json.dumps(result.__dict__)
+    assert len(_connections) == 1
+    executed_sql = " ".join(_connections[0].statement.lower().split())
+    assert "information_schema.columns" in executed_sql
+    assert "pg_catalog." in executed_sql
+    assert "not in ('pg_catalog', 'information_schema')" in executed_sql
+    assert not re.search(
+        r"\b(?:insert|update|delete|alter|drop|truncate|create)\b", executed_sql
+    )
+    assert _connections[0].parameters == {}
+
+
+def test_postgresql_incremental_resume_cancel_and_evidence_are_safe():
+    connector = PostgreSQLConnector(_provider)
+    initial = connector.snapshot(_request("snapshot"))
+
+    incremental = connector.incremental(
+        _request("incremental", checkpoint=initial.checkpoint)
+    )
+    resumed = connector.resume(_request("resume", checkpoint=initial.checkpoint))
+    cancelled = connector.cancel(_request("cancel", checkpoint=initial.checkpoint))
+    evidence = connector.evidence(_request("evidence"))
+
+    assert incremental.evidence["diff"]["changed"] is False
+    assert resumed.status == "succeeded"
+    assert cancelled.status == "cancelled"
+    assert evidence.evidence["secret_material"] is False
+
+
+def test_postgresql_cancel_probe_stops_catalog_collection():
+    with pytest.raises(ConnectorCancelledError):
+        PostgreSQLConnector(_provider).discover(
+            _request(cancel_probe=lambda: True)
+        )
+
+
+def test_postgresql_health_and_compatibility_fail_closed_without_driver(monkeypatch):
+    monkeypatch.setattr(
+        "importlib.util.find_spec",
+        lambda name: None if name == "psycopg2" else object(),
+    )
+
+    connector = PostgreSQLConnector(_provider)
+    health = connector.health({"credential_ref": "vault:postgresql/catalog"})
+    compatibility = connector.compatibility()
+
+    assert health.status == "degraded"
+    assert health.detail == "optional driver unavailable"
+    assert compatibility.compatible is False
+    assert compatibility.connector_version == "1.0.0"
+    assert compatibility.detail == "optional driver unavailable"