from __future__ import annotations import asyncio from src.jobs import asset_scanner_job class _Result: def scalar(self): return "00000000-0000-0000-0000-000000000001" def one(self): return 1, True class _Database: async def execute(self, statement, parameters=None): return _Result() class _Context: async def __aenter__(self): return _Database() async def __aexit__(self, exc_type, exc, traceback): return False def test_background_scanner_scopes_every_database_operation(monkeypatch) -> None: requested_projects: list[str | None] = [] def fake_db_context(project_id: str | None = None): requested_projects.append(project_id) return _Context() monkeypatch.setattr("src.db.base.get_db_context", fake_db_context) async def exercise_scanner_writes() -> None: run_id = await asset_scanner_job._start_discovery_run( "cron", ["k8s", "prometheus"], "shallow" ) assert run_id == "00000000-0000-0000-0000-000000000001" await asset_scanner_job._finish_discovery_run( run_id=run_id, status="success", total_assets=1, new_assets=1, modified_assets=0, duration_ms=10, error=None, ) inserted, modified = await asset_scanner_job._upsert_assets( [ { "asset_key": "host/test", "asset_type": "host", "host": "test", "namespace": None, "name": "test", "metadata": {}, "tags": ["source:test"], } ], run_id, ) assert (inserted, modified) == (1, 0) written = await asset_scanner_job._upsert_relationships( [ { "from_key": "host/test", "to_key": "monitoring/test", "relationship_type": "monitors", } ] ) assert written == 1 await asset_scanner_job._write_coverage_snapshots(run_id) await asset_scanner_job._log_aol_asset_discovered( run_id=run_id, triggered_by="cron", total=1, new=1, modified=0, duration_ms=10, status="success", error=None, scope=["k8s", "prometheus", "database_catalog"], scan_depth="shallow", ) asyncio.run(exercise_scanner_writes()) assert requested_projects == ["awoooi"] * 6 def test_database_catalog_collector_reads_metadata_only(monkeypatch) -> None: requested_projects: list[str | None] = [] class CatalogResult: def fetchall(self): return [ type( "CatalogRow", (), { "_mapping": { "database_name": "awoooi", "schema_name": "public", "relation_name": "incidents", "relation_kind": "table", "row_security_enabled": True, } }, )() ] class CatalogDatabase: async def execute(self, statement, parameters=None): return CatalogResult() class CatalogContext: async def __aenter__(self): return CatalogDatabase() async def __aexit__(self, exc_type, exc, traceback): return False def fake_db_context(project_id: str | None = None): requested_projects.append(project_id) return CatalogContext() monkeypatch.setattr("src.db.base.get_db_context", fake_db_context) assets, relationships = asyncio.run( asset_scanner_job._collect_database_catalog_assets() ) assert requested_projects == ["awoooi"] assert [asset["asset_type"] for asset in assets] == ["database", "table"] assert assets[0]["metadata"]["row_data_read"] is False assert assets[1]["metadata"] == { "database": "awoooi", "schema": "public", "relation_kind": "table", "row_security_enabled": True, "row_data_read": False, } assert relationships == [ { "from_key": "postgres/relation/awoooi/public/incidents", "to_key": "postgres/database/awoooi", "relationship_type": "depends_on", } ]