|
@@ -0,0 +1,180 @@
|
|
|
|
|
+from __future__ import annotations
|
|
|
|
|
+
|
|
|
|
|
+from contextlib import contextmanager
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class FakeResult:
|
|
|
|
|
+ def __init__(self, rows):
|
|
|
|
|
+ self.rows = rows
|
|
|
|
|
+
|
|
|
|
|
+ def mappings(self):
|
|
|
|
|
+ return self
|
|
|
|
|
+
|
|
|
|
|
+ def all(self):
|
|
|
|
|
+ return list(self.rows)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+class FakeConnection:
|
|
|
|
|
+ def __init__(self, rows):
|
|
|
|
|
+ self.rows = rows
|
|
|
|
|
+ self.calls = []
|
|
|
|
|
+
|
|
|
|
|
+ def execute(self, statement, parameters):
|
|
|
|
|
+ self.calls.append((str(statement), dict(parameters)))
|
|
|
|
|
+ return FakeResult(self.rows)
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def sample_rows():
|
|
|
|
|
+ return [
|
|
|
|
|
+ {
|
|
|
|
|
+ "schema_name": "public",
|
|
|
|
|
+ "asset_name": "orders",
|
|
|
|
|
+ "asset_type": "BASE TABLE",
|
|
|
|
|
+ "asset_comment": "orders",
|
|
|
|
|
+ "column_name": "id",
|
|
|
|
|
+ "ordinal_position": 1,
|
|
|
|
|
+ "data_type": "integer",
|
|
|
|
|
+ "is_nullable": "NO",
|
|
|
|
|
+ "column_default": None,
|
|
|
|
|
+ "column_comment": "primary key",
|
|
|
|
|
+ },
|
|
|
|
|
+ {
|
|
|
|
|
+ "schema_name": "public",
|
|
|
|
|
+ "asset_name": "orders",
|
|
|
|
|
+ "asset_type": "BASE TABLE",
|
|
|
|
|
+ "asset_comment": "orders",
|
|
|
|
|
+ "column_name": "amount",
|
|
|
|
|
+ "ordinal_position": 2,
|
|
|
|
|
+ "data_type": "numeric",
|
|
|
|
|
+ "is_nullable": "YES",
|
|
|
|
|
+ "column_default": None,
|
|
|
|
|
+ "column_comment": None,
|
|
|
|
|
+ },
|
|
|
|
|
+ ]
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def test_postgresql_collector_uses_allowlisted_catalog_and_bound_scope():
|
|
|
|
|
+ from app.core.data_research.catalog.models import CatalogScope
|
|
|
|
|
+ from app.core.data_research.catalog.postgresql import PostgreSqlCatalogCollector
|
|
|
|
|
+
|
|
|
|
|
+ connection = FakeConnection(sample_rows())
|
|
|
|
|
+ snapshot = PostgreSqlCatalogCollector().collect(
|
|
|
|
|
+ connection,
|
|
|
|
|
+ data_source_uid="source-1",
|
|
|
|
|
+ scope=CatalogScope(include_schemas=("public", "unsafe' OR 1=1 --")),
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ statement, parameters = connection.calls[0]
|
|
|
|
|
+ assert "information_schema.columns" in statement
|
|
|
|
|
+ assert "pg_catalog" in statement
|
|
|
|
|
+ assert "unsafe' OR 1=1 --" not in statement
|
|
|
|
|
+ assert parameters["include_schemas"] == ["public", "unsafe' OR 1=1 --"]
|
|
|
|
|
+ assert snapshot.assets[0].fields[1].key == "source-1:public.orders.amount"
|
|
|
|
|
+ assert snapshot.assets[0].fields[1].nullable is True
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def test_mysql_collector_uses_information_schema_and_bound_database():
|
|
|
|
|
+ from app.core.data_research.catalog.models import CatalogScope
|
|
|
|
|
+ from app.core.data_research.catalog.mysql import MySqlCatalogCollector
|
|
|
|
|
+
|
|
|
|
|
+ connection = FakeConnection(sample_rows())
|
|
|
|
|
+ snapshot = MySqlCatalogCollector().collect(
|
|
|
|
|
+ connection,
|
|
|
|
|
+ data_source_uid="source-2",
|
|
|
|
|
+ database_name="analytics' OR 1=1 --",
|
|
|
|
|
+ scope=CatalogScope(include_tables=("orders",)),
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ statement, parameters = connection.calls[0]
|
|
|
|
|
+ assert "information_schema.COLUMNS" in statement
|
|
|
|
|
+ assert "analytics' OR 1=1 --" not in statement
|
|
|
|
|
+ assert parameters["database_name"] == "analytics' OR 1=1 --"
|
|
|
|
|
+ assert parameters["include_tables"] == ["orders"]
|
|
|
|
|
+ assert snapshot.database_type == "mysql"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def test_catalog_service_uses_read_only_metadata_collection_purpose():
|
|
|
|
|
+ from app.core.data_research.catalog.models import CatalogScope
|
|
|
|
|
+ from app.core.data_research.catalog.postgresql import PostgreSqlCatalogCollector
|
|
|
|
|
+ from app.core.data_research.catalog.service import CatalogCollectionService
|
|
|
|
|
+
|
|
|
|
|
+ connection = FakeConnection(sample_rows())
|
|
|
|
|
+
|
|
|
|
|
+ class Manager:
|
|
|
|
|
+ def __init__(self):
|
|
|
|
|
+ self.calls = []
|
|
|
|
|
+
|
|
|
|
|
+ @contextmanager
|
|
|
|
|
+ def connect(self, uid, purpose):
|
|
|
|
|
+ self.calls.append((uid, purpose))
|
|
|
|
|
+ yield connection
|
|
|
|
|
+
|
|
|
|
|
+ manager = Manager()
|
|
|
|
|
+ service = CatalogCollectionService(
|
|
|
|
|
+ manager,
|
|
|
|
|
+ definition_resolver=lambda _uid: {
|
|
|
|
|
+ "database_type": "postgresql",
|
|
|
|
|
+ "database": "warehouse",
|
|
|
|
|
+ },
|
|
|
|
|
+ collector_resolver=lambda _database_type: PostgreSqlCatalogCollector(),
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ snapshot = service.collect("source-1", CatalogScope())
|
|
|
|
|
+
|
|
|
|
|
+ assert manager.calls == [("source-1", "metadata_collection")]
|
|
|
|
|
+ assert snapshot.assets[0].name == "orders"
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def _snapshot(fields):
|
|
|
|
|
+ from app.core.data_research.catalog.models import (
|
|
|
|
|
+ CatalogAsset,
|
|
|
|
|
+ CatalogField,
|
|
|
|
|
+ CatalogSnapshot,
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ return CatalogSnapshot(
|
|
|
|
|
+ data_source_uid="source-1",
|
|
|
|
|
+ database_type="postgresql",
|
|
|
|
|
+ assets=(
|
|
|
|
|
+ CatalogAsset(
|
|
|
|
|
+ key="source-1:public.orders",
|
|
|
|
|
+ schema="public",
|
|
|
|
|
+ name="orders",
|
|
|
|
|
+ asset_type="table",
|
|
|
|
|
+ fields=tuple(
|
|
|
|
|
+ CatalogField(
|
|
|
|
|
+ key=f"source-1:public.orders.{name}",
|
|
|
|
|
+ schema="public",
|
|
|
|
|
+ asset="orders",
|
|
|
|
|
+ name=name,
|
|
|
|
|
+ ordinal_position=ordinal,
|
|
|
|
|
+ data_type=data_type,
|
|
|
|
|
+ nullable=True,
|
|
|
|
|
+ )
|
|
|
|
|
+ for name, ordinal, data_type in fields
|
|
|
|
|
+ ),
|
|
|
|
|
+ ),
|
|
|
|
|
+ ),
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+
|
|
|
|
|
+def test_snapshot_diff_classifies_add_remove_type_change_and_rename():
|
|
|
|
|
+ from app.core.data_research.catalog.service import diff_snapshots
|
|
|
|
|
+
|
|
|
|
|
+ previous = _snapshot(
|
|
|
|
|
+ (("id", 1, "integer"), ("amount", 2, "numeric"), ("legacy", 3, "text"))
|
|
|
|
|
+ )
|
|
|
|
|
+ current = _snapshot(
|
|
|
|
|
+ (("id", 1, "bigint"), ("total", 2, "numeric"), ("created_at", 4, "timestamp"))
|
|
|
|
|
+ )
|
|
|
|
|
+
|
|
|
|
|
+ diff = diff_snapshots(previous, current)
|
|
|
|
|
+
|
|
|
|
|
+ assert [(item.before.name, item.after.name) for item in diff.renamed] == [
|
|
|
|
|
+ ("amount", "total")
|
|
|
|
|
+ ]
|
|
|
|
|
+ assert [(item.before.data_type, item.after.data_type) for item in diff.changed] == [
|
|
|
|
|
+ ("integer", "bigint")
|
|
|
|
|
+ ]
|
|
|
|
|
+ assert [item.name for item in diff.removed] == ["legacy"]
|
|
|
|
|
+ assert [item.name for item in diff.added] == ["created_at"]
|