diff --git a/api/changelog.d/compliance-overview-single-transaction.changed.md b/api/changelog.d/compliance-overview-single-transaction.changed.md new file mode 100644 index 00000000000..0e7a6714b3b --- /dev/null +++ b/api/changelog.d/compliance-overview-single-transaction.changed.md @@ -0,0 +1 @@ +Compliance overview ingest now runs in a single transaction per scan with a configurable `COPY` batch size (`DJANGO_COMPLIANCE_COPY_BATCH_SIZE`, default 2000), reducing write pressure on the database diff --git a/api/src/backend/tasks/jobs/scan.py b/api/src/backend/tasks/jobs/scan.py index 8290c45a7ab..b1f72ef76dc 100644 --- a/api/src/backend/tasks/jobs/scan.py +++ b/api/src/backend/tasks/jobs/scan.py @@ -6,7 +6,7 @@ import time import uuid from collections import defaultdict -from collections.abc import Iterable +from collections.abc import Callable, Iterable from datetime import UTC, datetime from typing import Any @@ -99,6 +99,11 @@ FINDINGS_MICRO_BATCH_SIZE = env.int("DJANGO_FINDINGS_MICRO_BATCH_SIZE", default=3000) # Controls how many rows each ORM bulk_create/bulk_update call sends to Postgres. SCAN_DB_BATCH_SIZE = env.int("DJANGO_SCAN_DB_BATCH_SIZE", default=1000) +# Rows per COPY statement when ingesting compliance requirement overviews. All +# batches of a scan share one transaction/commit; the batch size only bounds the +# client-side CSV buffer and how long each individual COPY statement runs on the +# writer (memory footprint, lock time and slow-statement logging under load). +COMPLIANCE_COPY_BATCH_SIZE = env.int("DJANGO_COMPLIANCE_COPY_BATCH_SIZE", default=2000) # Throttle scan progress persistence: minimum progress delta (fraction 0-1) # between two persisted progress updates. PROGRESS_THROTTLE_DELTA = env.float("DJANGO_SCAN_PROGRESS_THROTTLE_DELTA", default=0.01) @@ -356,21 +361,8 @@ def _bulk_update_resource_failed_findings_counts( raise -def _copy_compliance_requirement_rows( - tenant_id: str, rows: list[dict[str, Any]] -) -> None: - """Stream compliance requirement rows into Postgres using COPY. - - We leverage the admin connection (when available) to bypass the COPY + RLS - restriction, writing only the fields required by - ``ComplianceRequirementOverview``. - - Args: - tenant_id: Target tenant UUID. - rows: List of row dictionaries prepared by - :func:`create_compliance_requirements`. - """ - +def _compliance_requirement_rows_to_csv(rows: list[dict[str, Any]]) -> io.StringIO: + """Serialize compliance requirement rows into a CSV buffer for COPY.""" csv_buffer = io.StringIO() writer = csv.writer(csv_buffer) @@ -398,60 +390,86 @@ def _copy_compliance_requirement_rows( ) csv_buffer.seek(0) - copy_sql = ( - "COPY compliance_requirements_overviews (" - + ", ".join(COMPLIANCE_REQUIREMENT_COPY_COLUMNS) - + ") FROM STDIN WITH (FORMAT CSV, DELIMITER ',', QUOTE '\"', ESCAPE '\"', NULL '\\N')" - ) - - try: - with psycopg_connection(MainRouter.admin_db) as connection: - connection.autocommit = False - try: - with connection.cursor() as cursor: - cursor.execute(SET_CONFIG_QUERY, [POSTGRES_TENANT_VAR, tenant_id]) - cursor.copy_expert(copy_sql, csv_buffer) - connection.commit() - except Exception: - connection.rollback() - raise - finally: - csv_buffer.close() + return csv_buffer -def _persist_compliance_requirement_rows( - tenant_id: str, rows: Iterable[dict[str, Any]], batch_size: int = 10000 +def _copy_compliance_requirement_rows( + tenant_id: str, scan_id: str, rows: Iterable[dict[str, Any]], batch_size: int ) -> int: - """Persist compliance requirement rows using batched COPY with ORM fallback. + """Replace a scan's compliance requirement rows using batched COPY. - ``rows`` is consumed lazily in batches, so peak memory stays at ~``batch_size`` - rows instead of the full set. A batch that fails COPY falls back to an ORM - ``bulk_create`` of just that batch. + We leverage the admin connection (when available) to bypass the COPY + RLS + restriction. The scan's DELETE and every COPY batch run on one connection + inside a single transaction with a single commit, so the writer takes one + fsync per scan instead of one per batch, and a failed ingest rolls back + without committing a partial delete/insert (which a retry would otherwise + delete again, feeding dead rows to autovacuum). Args: tenant_id: Target tenant UUID. - rows: Iterable of row dictionaries reflecting the compliance overview - state for a scan. - batch_size: Number of rows per COPY batch (default: 10000). + scan_id: Scan whose previous rows are replaced. + rows: Iterable of row dictionaries, consumed lazily batch by batch. + batch_size: Number of rows per COPY statement. Returns: - int: total number of rows persisted. + int: total number of rows staged and committed. """ total_rows = 0 batch_num = 0 + copy_sql = ( + "COPY compliance_requirements_overviews (" + + ", ".join(COMPLIANCE_REQUIREMENT_COPY_COLUMNS) + + ") FROM STDIN WITH (FORMAT CSV, DELIMITER ',', QUOTE '\"', ESCAPE '\"', NULL '\\N')" + ) - for batch, _is_last in batched(rows, batch_size): - if not batch: - continue - batch_num += 1 + with psycopg_connection(MainRouter.admin_db) as connection: + connection.autocommit = False try: - _copy_compliance_requirement_rows(tenant_id, batch) - except Exception as error: - logger.exception( - f"COPY bulk insert for compliance requirements batch {batch_num} " - "failed; falling back to ORM bulk_create for this batch", - exc_info=error, - ) + with connection.cursor() as cursor: + cursor.execute(SET_CONFIG_QUERY, [POSTGRES_TENANT_VAR, tenant_id]) + # Idempotent re-run: clearing this scan's rows inside the same + # transaction keeps delete + reinsert atomic. + cursor.execute( + "DELETE FROM compliance_requirements_overviews " + "WHERE tenant_id = %s AND scan_id = %s", + [tenant_id, scan_id], + ) + for batch, _is_last in batched(rows, batch_size): + if not batch: + continue + batch_num += 1 + csv_buffer = _compliance_requirement_rows_to_csv(batch) + try: + cursor.copy_expert(copy_sql, csv_buffer) + finally: + csv_buffer.close() + total_rows += len(batch) + logger.info( + f"Compliance COPY batch {batch_num}: staged {len(batch)} rows " + f"({total_rows} total)" + ) + connection.commit() + except Exception: + connection.rollback() + raise + + return total_rows + + +def _bulk_create_compliance_requirement_rows( + tenant_id: str, scan_id: str, rows: Iterable[dict[str, Any]], batch_size: int +) -> int: + """Replace a scan's compliance requirement rows via the ORM. + + Fallback for when COPY is unavailable; the delete and every ``bulk_create`` + share one RLS transaction so the replacement stays atomic. + """ + total_rows = 0 + with rls_transaction(tenant_id): + ComplianceRequirementOverview.objects.filter(scan_id=scan_id).delete() + for batch, _is_last in batched(rows, batch_size): + if not batch: + continue fallback_objects = [ ComplianceRequirementOverview( id=row["id"], @@ -473,18 +491,52 @@ def _persist_compliance_requirement_rows( ) for row in batch ] - with rls_transaction(tenant_id): - ComplianceRequirementOverview.objects.bulk_create( - fallback_objects, batch_size=500 - ) + ComplianceRequirementOverview.objects.bulk_create( + fallback_objects, batch_size=500 + ) + total_rows += len(batch) + return total_rows - total_rows += len(batch) - logger.info( - f"Compliance COPY batch {batch_num}: inserted {len(batch)} rows " - f"({total_rows} total)" - ) - return total_rows +def _persist_compliance_requirement_rows( + tenant_id: str, + scan_id: str, + rows_factory: Callable[[], Iterable[dict[str, Any]]], + batch_size: int | None = None, +) -> int: + """Persist a scan's compliance requirement rows, replacing any previous ones. + + ``rows_factory`` must return a fresh row iterator on every call: the COPY + path consumes it lazily in batches (peak memory ~``batch_size`` rows), and + if COPY fails the whole ingest falls back to a single ORM transaction that + re-iterates the rows. + + Args: + tenant_id: Target tenant UUID. + scan_id: Scan whose compliance overview rows are being replaced. + rows_factory: Callable returning an iterable of row dictionaries. + batch_size: Rows per COPY/bulk_create batch (default: + ``COMPLIANCE_COPY_BATCH_SIZE``). + + Returns: + int: total number of rows persisted. + """ + if batch_size is None: + batch_size = COMPLIANCE_COPY_BATCH_SIZE + + try: + return _copy_compliance_requirement_rows( + tenant_id, scan_id, rows_factory(), batch_size + ) + except Exception as error: + logger.exception( + "COPY bulk insert for compliance requirements failed; " + "falling back to ORM bulk_create", + exc_info=error, + ) + return _bulk_create_compliance_requirement_rows( + tenant_id, scan_id, rows_factory(), batch_size + ) def _create_compliance_summaries( @@ -885,15 +937,19 @@ def _process_finding_micro_batch( # Denormalized resource arrays populated directly on insert # (was previously a separate bulk_update; saves a CASE WHEN # over thousands of rows per micro-batch). - resource_regions=[resource_instance.region] - if resource_instance.region - else [], - resource_services=[resource_instance.service] - if resource_instance.service - else [], - resource_types=[resource_instance.type] - if resource_instance.type - else [], + resource_regions=( + [resource_instance.region] + if resource_instance.region + else [] + ), + resource_services=( + [resource_instance.service] + if resource_instance.service + else [] + ), + resource_types=( + [resource_instance.type] if resource_instance.type else [] + ), ) findings_to_create.append(finding_instance) resource_denormalized_data.append( @@ -1699,8 +1755,10 @@ def create_compliance_requirements(tenant_id: str, scan_id: str): ) # Yield rows lazily (consumed batch-by-batch by COPY) so peak memory - # stays bounded; tally requirement_statuses in the same pass. + # stays bounded; tally requirement_statuses in the same pass. The + # ORM fallback re-iterates from scratch, so the tally resets first. def _iter_compliance_requirement_rows(): + requirement_statuses.clear() for region in regions: region_stats = region_requirement_stats.get(region, {}) region_findings = findings_count_by_compliance.get(region, {}) @@ -1764,12 +1822,10 @@ def _iter_compliance_requirement_rows(): "total_findings": total_findings, } - # Idempotent re-run: clear this scan's rows before re-inserting. - with rls_transaction(tenant_id): - ComplianceRequirementOverview.objects.filter(scan_id=scan_id).delete() - + # The delete of the scan's previous rows happens inside the same + # transaction as the inserts (see _copy_compliance_requirement_rows). requirements_created = _persist_compliance_requirement_rows( - tenant_id, _iter_compliance_requirement_rows() + tenant_id_str, scan_id_str, _iter_compliance_requirement_rows ) # Create pre-aggregated summaries for fast compliance overview lookups diff --git a/api/src/backend/tasks/tests/test_scan.py b/api/src/backend/tasks/tests/test_scan.py index 2fd2dde05f7..417f1ac3a27 100644 --- a/api/src/backend/tasks/tests/test_scan.py +++ b/api/src/backend/tasks/tests/test_scan.py @@ -2314,9 +2314,9 @@ def test_create_compliance_requirements_check_status_priority( create_compliance_requirements(tenant_id, scan_id) mock_persist.assert_called_once() - persisted_rows = mock_persist.call_args[0][1] + rows_factory = mock_persist.call_args[0][2] requirement_row = next( - row for row in persisted_rows if row["requirement_id"] == "1.1" + row for row in rows_factory() if row["requirement_id"] == "1.1" ) assert requirement_row["requirement_status"] == "FAIL" @@ -2454,18 +2454,26 @@ def copy_side_effect(sql, file_obj): } with patch.object(MainRouter, "admin_db", "admin"): - _copy_compliance_requirement_rows(str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows( + str(row["tenant_id"]), str(row["scan_id"]), [row], 2000 + ) mock_psycopg_connection.assert_called_once_with("admin") connection.cursor.assert_called_once() - cursor.execute.assert_called_once() + # One execute for set_config plus one for the scan's DELETE. + assert cursor.execute.call_count == 2 + delete_sql, delete_params = cursor.execute.call_args_list[1][0] + assert "DELETE FROM compliance_requirements_overviews" in delete_sql + assert delete_params == [str(row["tenant_id"]), str(row["scan_id"])] cursor.copy_expert.assert_called_once() + connection.commit.assert_called_once() csv_rows = list(csv.reader(StringIO(captured["data"]))) assert csv_rows[0][0] == str(row["id"]) assert csv_rows[0][5] == "" assert csv_rows[0][-1] == str(row["scan_id"]) + @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.filter") @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.bulk_create") @patch("tasks.jobs.scan.rls_transaction") @patch( @@ -2473,7 +2481,7 @@ def copy_side_effect(sql, file_obj): side_effect=Exception("copy failed"), ) def test_persist_compliance_requirement_rows_fallback( - self, mock_copy, mock_rls_transaction, mock_bulk_create + self, mock_copy, mock_rls_transaction, mock_bulk_create, mock_filter ): inserted_at = datetime.now(UTC) row = { @@ -2494,16 +2502,22 @@ def test_persist_compliance_requirement_rows_fallback( } tenant_id = row["tenant_id"] + scan_id = str(row["scan_id"]) ctx = MagicMock() ctx.__enter__.return_value = None ctx.__exit__.return_value = False mock_rls_transaction.return_value = ctx - _persist_compliance_requirement_rows(tenant_id, [row]) + _persist_compliance_requirement_rows(tenant_id, scan_id, lambda: [row]) - mock_copy.assert_called_once_with(tenant_id, [row]) + mock_copy.assert_called_once() + assert mock_copy.call_args[0][0] == tenant_id + assert mock_copy.call_args[0][1] == scan_id mock_rls_transaction.assert_called_once_with(tenant_id) + # The fallback replaces the scan's rows: delete + insert atomically. + mock_filter.assert_called_once_with(scan_id=scan_id) + mock_filter.return_value.delete.assert_called_once() mock_bulk_create.assert_called_once() args, kwargs = mock_bulk_create.call_args @@ -2515,13 +2529,18 @@ def test_persist_compliance_requirement_rows_fallback( @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.bulk_create") @patch("tasks.jobs.scan.rls_transaction") - @patch("tasks.jobs.scan._copy_compliance_requirement_rows") + @patch("tasks.jobs.scan._copy_compliance_requirement_rows", return_value=0) def test_persist_compliance_requirement_rows_no_rows( self, mock_copy, mock_rls_transaction, mock_bulk_create ): - _persist_compliance_requirement_rows(str(uuid.uuid4()), []) + # Even with no rows the COPY path runs: it must clear the scan's + # previous rows so a re-run with fewer findings drops stale data. + total = _persist_compliance_requirement_rows( + str(uuid.uuid4()), str(uuid.uuid4()), lambda: [] + ) - mock_copy.assert_not_called() + assert total == 0 + mock_copy.assert_called_once() mock_rls_transaction.assert_not_called() mock_bulk_create.assert_not_called() @@ -2610,11 +2629,12 @@ def copy_side_effect(sql, file_obj): ] with patch.object(MainRouter, "admin_db", "admin"): - _copy_compliance_requirement_rows(tenant_id, rows) + _copy_compliance_requirement_rows(tenant_id, str(scan_id), rows, 2000) mock_psycopg_connection.assert_called_once_with("admin") connection.cursor.assert_called_once() - cursor.execute.assert_called_once() + # set_config + DELETE of the scan's previous rows. + assert cursor.execute.call_count == 2 cursor.copy_expert.assert_called_once() csv_rows = list(csv.reader(StringIO(captured["data"]))) @@ -2644,6 +2664,60 @@ def copy_side_effect(sql, file_obj): assert csv_rows[2][5] == "2.0" assert csv_rows[2][9] == "MANUAL" + @patch("tasks.jobs.scan.psycopg_connection") + def test_copy_compliance_requirement_rows_batches_share_one_transaction( + self, mock_psycopg_connection, settings + ): + """Every COPY batch runs on the same connection with a single commit.""" + settings.DATABASES.setdefault("admin", settings.DATABASES["default"]) + + connection = MagicMock() + cursor = MagicMock() + cursor_context = MagicMock() + cursor_context.__enter__.return_value = cursor + cursor_context.__exit__.return_value = False + connection.cursor.return_value = cursor_context + connection.__enter__.return_value = connection + connection.__exit__.return_value = False + + context_manager = MagicMock() + context_manager.__enter__.return_value = connection + context_manager.__exit__.return_value = False + mock_psycopg_connection.return_value = context_manager + + tenant_id = str(uuid.uuid4()) + scan_id = str(uuid.uuid4()) + inserted_at = datetime.now(UTC) + rows = [ + { + "id": uuid.uuid4(), + "tenant_id": tenant_id, + "inserted_at": inserted_at, + "compliance_id": "cisa_aws", + "framework": "CISA", + "version": "1.0", + "description": f"Requirement {index}", + "region": "us-east-1", + "requirement_id": f"req-{index}", + "requirement_status": "PASS", + "passed_checks": 1, + "failed_checks": 0, + "total_checks": 1, + "scan_id": scan_id, + } + for index in range(3) + ] + + with patch.object(MainRouter, "admin_db", "admin"): + total = _copy_compliance_requirement_rows(tenant_id, scan_id, rows, 1) + + assert total == 3 + # One connection, three COPY statements, one commit for the whole scan. + mock_psycopg_connection.assert_called_once_with("admin") + assert cursor.copy_expert.call_count == 3 + connection.commit.assert_called_once() + connection.rollback.assert_not_called() + @patch("tasks.jobs.scan.psycopg_connection") def test_copy_compliance_requirement_rows_null_values( self, mock_psycopg_connection, settings @@ -2691,7 +2765,9 @@ def copy_side_effect(sql, file_obj): } with patch.object(MainRouter, "admin_db", "admin"): - _copy_compliance_requirement_rows(str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows( + str(row["tenant_id"]), str(row["scan_id"]), [row], 2000 + ) csv_rows = list(csv.reader(StringIO(captured["data"]))) assert len(csv_rows) == 1 @@ -2747,7 +2823,9 @@ def copy_side_effect(sql, file_obj): } with patch.object(MainRouter, "admin_db", "admin"): - _copy_compliance_requirement_rows(str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows( + str(row["tenant_id"]), str(row["scan_id"]), [row], 2000 + ) # Verify CSV was generated (csv module handles escaping automatically) csv_rows = list(csv.reader(StringIO(captured["data"]))) @@ -2808,7 +2886,9 @@ def copy_side_effect(sql, file_obj): before_call = datetime.now(UTC) with patch.object(MainRouter, "admin_db", "admin"): - _copy_compliance_requirement_rows(str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows( + str(row["tenant_id"]), str(row["scan_id"]), [row], 2000 + ) after_call = datetime.now(UTC) csv_rows = list(csv.reader(StringIO(captured["data"]))) @@ -2861,7 +2941,9 @@ def test_copy_compliance_requirement_rows_transaction_rollback_on_copy_error( with patch.object(MainRouter, "admin_db", "admin"): with pytest.raises(Exception, match="COPY command failed"): - _copy_compliance_requirement_rows(str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows( + str(row["tenant_id"]), str(row["scan_id"]), [row], 2000 + ) # Verify rollback was called connection.rollback.assert_called_once() @@ -2909,7 +2991,9 @@ def test_copy_compliance_requirement_rows_transaction_rollback_on_set_config_err with patch.object(MainRouter, "admin_db", "admin"): with pytest.raises(Exception, match="SET prowler.tenant_id failed"): - _copy_compliance_requirement_rows(str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows( + str(row["tenant_id"]), str(row["scan_id"]), [row], 2000 + ) # Verify rollback was called connection.rollback.assert_called_once() @@ -2955,7 +3039,9 @@ def test_copy_compliance_requirement_rows_commit_on_success( } with patch.object(MainRouter, "admin_db", "admin"): - _copy_compliance_requirement_rows(str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows( + str(row["tenant_id"]), str(row["scan_id"]), [row], 2000 + ) # Verify commit was called and rollback was not connection.commit.assert_called_once() @@ -2966,9 +3052,10 @@ def test_copy_compliance_requirement_rows_commit_on_success( @patch("tasks.jobs.scan._copy_compliance_requirement_rows") def test_persist_compliance_requirement_rows_success(self, mock_copy): """Test successful COPY path without fallback to ORM.""" - mock_copy.return_value = None # Success, no exception + mock_copy.return_value = 1 # Success, no exception tenant_id = str(uuid.uuid4()) + scan_id = str(uuid.uuid4()) rows = [ { "id": uuid.uuid4(), @@ -2984,16 +3071,21 @@ def test_persist_compliance_requirement_rows_success(self, mock_copy): "passed_checks": 1, "failed_checks": 0, "total_checks": 1, - "scan_id": uuid.uuid4(), + "scan_id": scan_id, } ] - _persist_compliance_requirement_rows(tenant_id, rows) + total = _persist_compliance_requirement_rows(tenant_id, scan_id, lambda: rows) - # Verify COPY was called - mock_copy.assert_called_once_with(tenant_id, rows) + assert total == 1 + mock_copy.assert_called_once() + copy_args = mock_copy.call_args[0] + assert copy_args[0] == tenant_id + assert copy_args[1] == scan_id + assert list(copy_args[2]) == rows @patch("tasks.jobs.scan.logger") + @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.filter") @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.bulk_create") @patch("tasks.jobs.scan.rls_transaction") @patch( @@ -3001,7 +3093,12 @@ def test_persist_compliance_requirement_rows_success(self, mock_copy): side_effect=Exception("COPY failed"), ) def test_persist_compliance_requirement_rows_fallback_logging( - self, mock_copy, mock_rls_transaction, mock_bulk_create, mock_logger + self, + mock_copy, + mock_rls_transaction, + mock_bulk_create, + mock_filter, + mock_logger, ): """Test logger.exception is called when COPY fails and fallback occurs.""" tenant_id = str(uuid.uuid4()) @@ -3027,7 +3124,9 @@ def test_persist_compliance_requirement_rows_fallback_logging( ctx.__exit__.return_value = False mock_rls_transaction.return_value = ctx - _persist_compliance_requirement_rows(tenant_id, [row]) + _persist_compliance_requirement_rows( + tenant_id, str(row["scan_id"]), lambda: [row] + ) # Verify logger.exception was called mock_logger.exception.assert_called_once() @@ -3036,6 +3135,7 @@ def test_persist_compliance_requirement_rows_fallback_logging( assert "falling back to ORM" in args[0] assert kwargs.get("exc_info") is not None + @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.filter") @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.bulk_create") @patch("tasks.jobs.scan.rls_transaction") @patch( @@ -3043,7 +3143,7 @@ def test_persist_compliance_requirement_rows_fallback_logging( side_effect=Exception("copy failed"), ) def test_persist_compliance_requirement_rows_fallback_multiple_rows( - self, mock_copy, mock_rls_transaction, mock_bulk_create + self, mock_copy, mock_rls_transaction, mock_bulk_create, mock_filter ): """Test ORM fallback with multiple rows.""" tenant_id = str(uuid.uuid4()) @@ -3090,10 +3190,14 @@ def test_persist_compliance_requirement_rows_fallback_multiple_rows( ctx.__exit__.return_value = False mock_rls_transaction.return_value = ctx - _persist_compliance_requirement_rows(tenant_id, rows) + total = _persist_compliance_requirement_rows( + tenant_id, str(scan_id), lambda: rows + ) - mock_copy.assert_called_once_with(tenant_id, rows) + assert total == 2 + mock_copy.assert_called_once() mock_rls_transaction.assert_called_once_with(tenant_id) + mock_filter.assert_called_once_with(scan_id=str(scan_id)) mock_bulk_create.assert_called_once() args, kwargs = mock_bulk_create.call_args @@ -3117,6 +3221,7 @@ def test_persist_compliance_requirement_rows_fallback_multiple_rows( assert objects[1].passed_checks == 2 assert objects[1].failed_checks == 3 + @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.filter") @patch("tasks.jobs.scan.ComplianceRequirementOverview.objects.bulk_create") @patch("tasks.jobs.scan.rls_transaction") @patch( @@ -3124,7 +3229,7 @@ def test_persist_compliance_requirement_rows_fallback_multiple_rows( side_effect=Exception("copy failed"), ) def test_persist_compliance_requirement_rows_fallback_all_fields( - self, mock_copy, mock_rls_transaction, mock_bulk_create + self, mock_copy, mock_rls_transaction, mock_bulk_create, mock_filter ): """Test ORM fallback correctly maps all fields from row dict to model.""" tenant_id = str(uuid.uuid4()) @@ -3154,7 +3259,7 @@ def test_persist_compliance_requirement_rows_fallback_all_fields( ctx.__exit__.return_value = False mock_rls_transaction.return_value = ctx - _persist_compliance_requirement_rows(tenant_id, [row]) + _persist_compliance_requirement_rows(tenant_id, str(scan_id), lambda: [row]) args, kwargs = mock_bulk_create.call_args objects = args[0]