|
@@ -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"
|