diff --git a/apps/api/src/services/iwooos_security_asset_control_plane.py b/apps/api/src/services/iwooos_security_asset_control_plane.py index c7f61ff2b..816290061 100644 --- a/apps/api/src/services/iwooos_security_asset_control_plane.py +++ b/apps/api/src/services/iwooos_security_asset_control_plane.py @@ -15,7 +15,8 @@ from src.db.base import get_db_context logger = structlog.get_logger(__name__) _SCHEMA_VERSION = "iwooos_security_asset_control_plane_v1" -_QUERY_TIMEOUT_SECONDS = 5.0 +_QUERY_TIMEOUT_SECONDS = 6.0 +_SOURCE_QUERY_TIMEOUT_SECONDS = 4.0 _DISCOVERY_FRESHNESS_SECONDS = 2 * 60 * 60 _WINDOW_HOURS = 24 @@ -1272,30 +1273,45 @@ def build_unavailable_iwooos_security_asset_control_plane( async def _read_public_source( - db: Any, *, + project_id: str, source_id: str, statement: Any, result_mode: str, default: Any, ) -> tuple[Any, dict[str, Any]]: """Isolate one aggregate query so an optional source cannot abort the readback.""" - try: - async with db.begin_nested(): + + async def execute_source() -> Any: + async with get_db_context(project_id) as db: result = await db.execute(statement) if result_mode == "all": - value = result.fetchall() - elif result_mode == "one": - value = result.one() - elif result_mode == "one_or_none": - value = result.one_or_none() - else: - raise ValueError("unsupported_result_mode") + return result.fetchall() + if result_mode == "one": + return result.one() + if result_mode == "one_or_none": + return result.one_or_none() + raise ValueError("unsupported_result_mode") + + try: + value = await asyncio.wait_for( + execute_source(), timeout=_SOURCE_QUERY_TIMEOUT_SECONDS + ) return value, { "source_id": source_id, "status": "ready", "reason_code": None, } + except TimeoutError: + logger.warning( + "iwooos_security_asset_control_plane_source_timeout", + source_id=source_id, + ) + return default, { + "source_id": source_id, + "status": "unavailable", + "reason_code": f"{source_id}_query_timeout", + } except Exception as exc: # noqa: BLE001 - isolate and redact source failures logger.warning( "iwooos_security_asset_control_plane_source_unavailable", @@ -1332,11 +1348,9 @@ class IwoooSSecurityAssetControlPlaneService: ) async def _load_snapshot(self) -> dict[str, Any]: - async with get_db_context("awoooi") as db: - source_health: list[dict[str, Any]] = [] - - discovery_row, health = await _read_public_source( - db, + source_results = await asyncio.gather( + _read_public_source( + project_id="awoooi", source_id="asset_discovery", statement=_sql( """ @@ -1369,11 +1383,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="one_or_none", default=None, - ) - source_health.append(health) - - inventory_rows, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="asset_inventory", statement=_sql( """ @@ -1394,11 +1406,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="all", default=[], - ) - source_health.append(health) - - coverage_rows, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="automation_coverage", statement=_sql( """ @@ -1417,11 +1427,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="all", default=[], - ) - source_health.append(health) - - relationship_row, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="asset_relationship", statement=_sql( """ @@ -1441,11 +1449,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="one", default={}, - ) - source_health.append(health) - - compliance_rows, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="security_compliance", statement=_sql( """ @@ -1463,11 +1469,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="all", default=[], - ) - source_health.append(health) - - automation_rows, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="automation_operations", statement=_sql( """ @@ -1480,11 +1484,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="all", default=[], - ) - source_health.append(health) - - ai_trace_row, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="ai_decision_trace", statement=_sql( """ @@ -1498,11 +1500,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="one", default={}, - ) - source_health.append(health) - - siem_row, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="siem_runtime", statement=_sql( """ @@ -1548,11 +1548,9 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="one", default={}, - ) - source_health.append(health) - - audit_row, health = await _read_public_source( - db, + ), + _read_public_source( + project_id="awoooi", source_id="audit_runtime", statement=_sql( """ @@ -1584,8 +1582,31 @@ class IwoooSSecurityAssetControlPlaneService: ), result_mode="one", default={}, - ) - source_health.append(health) + ), + ) + + ( + (discovery_row, discovery_health), + (inventory_rows, inventory_health), + (coverage_rows, coverage_health), + (relationship_row, relationship_health), + (compliance_rows, compliance_health), + (automation_rows, automation_health), + (ai_trace_row, ai_trace_health), + (siem_row, siem_health), + (audit_row, audit_health), + ) = source_results + source_health = [ + discovery_health, + inventory_health, + coverage_health, + relationship_health, + compliance_health, + automation_health, + ai_trace_health, + siem_health, + audit_health, + ] return build_iwooos_security_asset_control_plane( discovery_row=discovery_row, diff --git a/apps/api/tests/test_iwooos_security_asset_control_plane.py b/apps/api/tests/test_iwooos_security_asset_control_plane.py index 85d153d29..f4f2c8946 100644 --- a/apps/api/tests/test_iwooos_security_asset_control_plane.py +++ b/apps/api/tests/test_iwooos_security_asset_control_plane.py @@ -224,24 +224,28 @@ def test_partial_source_failure_is_visible_without_erasing_core_inventory() -> N assert payload["work_items"][0]["work_item_id"] == "AIA-P0-006-00" -def test_read_public_source_returns_redacted_reason_code_on_query_failure() -> None: - class NestedContext: +def test_read_public_source_returns_redacted_reason_code_on_query_failure( + monkeypatch, +) -> None: + class FailingDb: + async def execute(self, statement): + raise RuntimeError("private table detail") + + class DbContext: async def __aenter__(self): - return None + return FailingDb() async def __aexit__(self, exc_type, exc, traceback): return False - class FailingDb: - def begin_nested(self): - return NestedContext() - - async def execute(self, statement): - raise RuntimeError("private table detail") + monkeypatch.setattr( + "src.services.iwooos_security_asset_control_plane.get_db_context", + lambda project_id: DbContext(), + ) value, health = asyncio.run( _read_public_source( - FailingDb(), + project_id="awoooi", source_id="audit_runtime", statement="SELECT 1", result_mode="one", @@ -307,9 +311,51 @@ def test_service_uses_canonical_awoooi_project_scope(monkeypatch) -> None: ) payload = asyncio.run(service.get_snapshot()) - assert requested_projects == ["awoooi"] + assert requested_projects == ["awoooi"] * 9 assert payload["status"] == "degraded" - assert payload["reason_code"] == "database_query_failed" + assert payload["source_status"] == "live_database_unavailable" + assert payload["summary"]["unavailable_source_count"] == 9 + + +def test_service_reads_independent_sources_concurrently(monkeypatch) -> None: + service = IwoooSSecurityAssetControlPlaneService() + concurrency = {"active": 0, "maximum": 0} + + class EmptyResult: + def fetchall(self): + return [] + + def one(self): + return {} + + def one_or_none(self): + return None + + class SlowDb: + async def execute(self, statement): + concurrency["active"] += 1 + concurrency["maximum"] = max(concurrency["maximum"], concurrency["active"]) + await asyncio.sleep(0.02) + concurrency["active"] -= 1 + return EmptyResult() + + class DbContext: + async def __aenter__(self): + return SlowDb() + + async def __aexit__(self, exc_type, exc, traceback): + return False + + monkeypatch.setattr( + "src.services.iwooos_security_asset_control_plane.get_db_context", + lambda project_id: DbContext(), + ) + payload = asyncio.run(service._load_snapshot()) + + assert concurrency["maximum"] == 9 + assert payload["summary"]["source_count"] == 9 + assert payload["summary"]["ready_source_count"] == 9 + assert payload["summary"]["unavailable_source_count"] == 0 def test_public_api_returns_live_aggregate(monkeypatch) -> None: