DataOps-platform/tests/test_postgresql_enterprise_connector.py

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"