perf(api): harden compliance-overview ingest & partition table#11870
perf(api): harden compliance-overview ingest & partition table#11870pedrooot wants to merge 3 commits into
Conversation
github-actions
Bot
added
component/api
review-django-migrations
|
✅ All necessary |
|
✅ No Conflicts No conflict markers, and the branch merges cleanly into its base. |
📝 WalkthroughWalkthroughThis PR converts ChangesPartitioning and hardened ingest
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
🔒 Container Security ScanImage: ✅ No Vulnerabilities DetectedThe container image passed all security checks. No known CVEs were found.📋 Resources:
|
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #11870 +/- ##
==========================================
- Coverage 93.83% 93.81% -0.02%
==========================================
Files 263 264 +1
Lines 38692 38701 +9
==========================================
+ Hits 36305 36306 +1
- Misses 2387 2395 +8
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
api/src/backend/tasks/tests/test_scan.py (1)
2887-2923: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMissing test coverage for multi-batch connection reuse/throttling.
_persist_compliance_requirement_rowsreuses a single admin connection across batches and supportsthrottle_seconds, but every test here only exercises a single batch. There's no assertion thatpsycopg_connectionis opened exactly once across multiple batches, or that throttling is applied between batches — both are core claims of this hardening work per the changelog entry.Consider adding a test with
batch_sizesmaller than the row count to assertmock_psycopg_connection.assert_called_once()across batches, plus a throttle test assertingtime.sleepis invoked between batches.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@api/src/backend/tasks/tests/test_scan.py` around lines 2887 - 2923, The current `_persist_compliance_requirement_rows` test only covers a single successful batch and does not verify the new multi-batch behavior. Add a test that calls `_persist_compliance_requirement_rows` with more rows than the configured batch size so it processes multiple batches, then assert `psycopg_connection` is opened only once and `_copy_compliance_requirement_rows` is invoked for each batch. Also add a separate throttling test for the same function that patches `time.sleep` and verifies it is called between batches when `throttle_seconds` is set.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@api/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.py`:
- Line 139: The reverse migration in the partition cleanup logic is issuing an
invalid ALTER INDEX ... DETACH PARTITION statement, which will fail in
PostgreSQL. Update the migration rollback code around the partition/index
handling to remove the explicit index detach and only detach the table
partition; the attached index will be detached automatically. Use the
migration’s reverse SQL block that references INDEX and the _default partition
to make the change.
---
Outside diff comments:
In `@api/src/backend/tasks/tests/test_scan.py`:
- Around line 2887-2923: The current `_persist_compliance_requirement_rows` test
only covers a single successful batch and does not verify the new multi-batch
behavior. Add a test that calls `_persist_compliance_requirement_rows` with more
rows than the configured batch size so it processes multiple batches, then
assert `psycopg_connection` is opened only once and
`_copy_compliance_requirement_rows` is invoked for each batch. Also add a
separate throttling test for the same function that patches `time.sleep` and
verifies it is called between batches when `throttle_seconds` is set.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 8ccbe2a7-46fe-4392-a7eb-3628d4d70fdb
📒 Files selected for processing (7)
api/CHANGELOG.mdapi/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.pyapi/src/backend/api/models.pyapi/src/backend/api/partitions.pyapi/src/backend/config/settings/partitions.pyapi/src/backend/tasks/jobs/scan.pyapi/src/backend/tasks/tests/test_scan.py
| with cursor() as cur: | ||
| # Detach the default partition (keeps its own pkey/index/FKs) and drop the | ||
| # partitioned parent, then rename the default back to the original table. | ||
| cur.execute(f"ALTER INDEX {INDEX} DETACH PARTITION {INDEX}_default;") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the migration and nearby definitions
FILE="api/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.py"
echo "== file size =="
wc -l "$FILE"
echo "== outline (if available) =="
ast-grep outline "$FILE" --view expanded || true
echo "== relevant lines around reverse migration =="
sed -n '110,170p' "$FILE" | cat -n
echo "== search for RowLevelSecurityConstraint deconstruct / serialization =="
rg -n "class RowLevelSecurityConstraint|def deconstruct|partition_name|statements|DEFAULT_RLS|ADD_CONSTRAINT|AddConstraint" api/src/backend/api -SRepository: prowler-cloud/prowler
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== RowLevelSecurityConstraint in api/src/backend/api/rls.py =="
sed -n '1,220p' api/src/backend/api/rls.py | cat -n
echo
echo "== Migration state/database operations in api/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.py =="
sed -n '160,220p' api/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.py | cat -n
echo
echo "== Exact AddConstraint usages for this migration =="
rg -n "0098_partition_compliance_requirement_overviews.py|AddConstraint|RemoveConstraint|SeparateDatabaseAndState|RunPython|AlterModelOptions|AlterModelManagers|RemoveConstraint|AddIndex" api/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.py -n -SRepository: prowler-cloud/prowler
Length of output: 10787
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path("api/src/backend/api/rls.py")
text = p.read_text()
for needle in ["class RowLevelSecurityConstraint", "def deconstruct", "self.partition_name", "self.statements"]:
idx = text.find(needle)
print(f"\n--- {needle} ---")
if idx != -1:
start = max(0, text.rfind("\n", 0, idx-200))
end = min(len(text), idx+800)
print(text[start:end])
PYRepository: prowler-cloud/prowler
Length of output: 4369
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Locate ComplianceRequirementOverview model =="
rg -n "class ComplianceRequirementOverview|RowLevelSecurityConstraint\\(|partition_name=\"default\"|statements=\\[" api/src/backend/api/models.py api/src/backend/api/partitions.py api/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.py -n -S
echo
echo "== Model block around ComplianceRequirementOverview =="
python3 - <<'PY'
from pathlib import Path
text = Path("api/src/backend/api/models.py").read_text().splitlines()
# Print a window around the first occurrence of ComplianceRequirementOverview
for i, line in enumerate(text):
if "class ComplianceRequirementOverview" in line:
start = max(0, i - 20)
end = min(len(text), i + 120)
for j in range(start, end):
print(f"{j+1:5d}: {text[j]}")
break
PY
echo
echo "== Partition helpers around the relevant class =="
python3 - <<'PY'
from pathlib import Path
text = Path("api/src/backend/api/partitions.py").read_text().splitlines()
for i, line in enumerate(text):
if "rls_statements" in line or "class Partition" in line:
start = max(0, i - 25)
end = min(len(text), i + 120)
for j in range(start, end):
print(f"{j+1:5d}: {text[j]}")
print("\n---\n")
if i > 250:
break
PYRepository: prowler-cloud/prowler
Length of output: 50377
🌐 Web query:
PostgreSQL ALTER INDEX DETACH PARTITION documentation
💡 Result:
In PostgreSQL, there is no direct command such as ALTER INDEX DETACH PARTITION. Instead, the detachment of an index is a side effect of detaching a partition from a table. When you detach a partition from a partitioned table using the ALTER TABLE... DETACH PARTITION command, any indexes that were attached to the parent table's indexes are automatically detached from them [1][2]. The relevant command to perform this operation is: ALTER TABLE parent_table_name DETACH PARTITION partition_name; Key technical details: 1. Relationship: While ALTER INDEX ATTACH PARTITION exists to explicitly link an index on a partition to an index on the partitioned parent table [3][4], the corresponding detachment occurs automatically during the table-level partition detachment [1][2]. 2. Behavior: Once the parent partition is detached, the detached partition remains as a standalone table, and its indexes are no longer bound to the parent's indexes [1][5]. 3. Execution: You can execute this as a standard (blocking) detachment or, in supported versions, use the CONCURRENTLY option to reduce locking requirements while detaching [1][2]. Because ALTER INDEX is effectively an alias for forms of ALTER TABLE that apply to indexes [3][4], all partition-related structural changes, including the detachment of indexes that were previously attached to a parent index, are managed through the ALTER TABLE command interface.
Citations:
- 1: https://www.postgresql.org/docs/19/sql-altertable.html
- 2: https://www.postgresql.org/docs/18/sql-altertable.html
- 3: https://www.postgresql.org/docs/19/sql-alterindex.html
- 4: https://www.postgresql.org/docs/17/sql-alterindex.html
- 5: https://www.postgresql.org/docs/15/sql-altertable.html
Remove the invalid index detach on line 139
ALTER INDEX ... DETACH PARTITION isn’t valid PostgreSQL, so the reverse migration will fail. Detach the table partition only; its attached index is detached automatically.
🧰 Tools
🪛 OpenGrep (1.23.0)
[ERROR] 139-139: SQL query built via f-string passed to execute()/executemany(). Use parameterized queries with placeholders instead.
(coderabbit.sql-injection.python-fstring-execute)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In
`@api/src/backend/api/migrations/0098_partition_compliance_requirement_overviews.py`
at line 139, The reverse migration in the partition cleanup logic is issuing an
invalid ALTER INDEX ... DETACH PARTITION statement, which will fail in
PostgreSQL. Update the migration rollback code around the partition/index
handling to remove the explicit index detach and only detach the table
partition; the attached index will be detached automatically. Use the
migration’s reverse SQL block that references INDEX and the _default partition
to make the change.
Source: Linters/SAST tools
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
api/src/backend/tasks/tests/test_scan.py (2)
2444-2445: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove vestigial
MainRouter.admin_dbpatch.These tests now pass
connectionexplicitly to_copy_compliance_requirement_rows, which (per the current implementation) no longer looks upMainRouter.admin_db. Thewith patch.object(MainRouter, "admin_db", "admin"):wrapper is leftover from the pre-refactor design and is dead weight in each of these tests.🧹 Example cleanup (repeat for each occurrence)
- with patch.object(MainRouter, "admin_db", "admin"): - _copy_compliance_requirement_rows(connection, str(row["tenant_id"]), [row]) + _copy_compliance_requirement_rows(connection, str(row["tenant_id"]), [row])Also applies to: 2592-2593, 2660-2661, 2704-2705, 2753-2754, 2793-2797, 2833-2837, 2870-2871
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@api/src/backend/tasks/tests/test_scan.py` around lines 2444 - 2445, Remove the leftover MainRouter.admin_db patch in the tests around _copy_compliance_requirement_rows, since the helper now uses the explicitly passed connection and no longer depends on MainRouter.admin_db. Update each affected test case to call _copy_compliance_requirement_rows directly without the with patch.object(MainRouter, "admin_db", "admin") wrapper, keeping the existing connection and row setup unchanged.
2412-2870: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate connection/cursor mock boilerplate across tests.
The
connection = MagicMock(); cursor = MagicMock(); cursor_context = MagicMock(); ...; connection.cursor.return_value = cursor_contextsetup is repeated verbatim in at least 8 test methods. A pytest fixture returning(connection, cursor)would remove this duplication and make future signature changes easier to maintain.♻️ Example fixture
`@pytest.fixture` def mock_copy_connection(): 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 return connection, cursor🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@api/src/backend/tasks/tests/test_scan.py` around lines 2412 - 2870, The connection/cursor setup is duplicated across multiple compliance COPY tests, making the suite noisy and harder to maintain. Extract the repeated MagicMock context-manager wiring into a shared pytest fixture (for example, one returning the connection and cursor used by tests like test_copy_compliance_requirement_rows_streams_csv and test_copy_compliance_requirement_rows_multiple_rows), then update each test to use that fixture instead of recreating the same boilerplate inline. Keep the existing cursor.copy_expert and cursor.execute assertions unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@api/src/backend/tasks/tests/test_scan.py`:
- Around line 2444-2445: Remove the leftover MainRouter.admin_db patch in the
tests around _copy_compliance_requirement_rows, since the helper now uses the
explicitly passed connection and no longer depends on MainRouter.admin_db.
Update each affected test case to call _copy_compliance_requirement_rows
directly without the with patch.object(MainRouter, "admin_db", "admin") wrapper,
keeping the existing connection and row setup unchanged.
- Around line 2412-2870: The connection/cursor setup is duplicated across
multiple compliance COPY tests, making the suite noisy and harder to maintain.
Extract the repeated MagicMock context-manager wiring into a shared pytest
fixture (for example, one returning the connection and cursor used by tests like
test_copy_compliance_requirement_rows_streams_csv and
test_copy_compliance_requirement_rows_multiple_rows), then update each test to
use that fixture instead of recreating the same boilerplate inline. Keep the
existing cursor.copy_expert and cursor.execute assertions unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: c84382ac-d551-44cb-b3f9-2ccc74015526
📒 Files selected for processing (1)
api/src/backend/tasks/tests/test_scan.py
|
closed by #11875 |
1 participant
Description
On every scan we delete the scan's existing rows and re-insert the whole set, streaming them into
compliance_requirements_overviewswith back-to-backCOPYstatements that take 13–15s each. On a plain, ever-growing table that means constant commit/fsync waits, lock contention, and a lot of autovacuum work chasing the dead tuples the delete-then-reinsert leaves behind. Under load it's enough to tip the writer over.What changed
Ingest (
tasks/jobs/scan.py). The COPY batch size is now configurable viaDJANGO_COMPLIANCE_COPY_BATCH_SIZE(default dropped from 10000 to 5000), with an optional pause between batches viaDJANGO_COMPLIANCE_COPY_THROTTLE_SECONDS(default 0, so behaviour is unchanged unless someone turns it up). Smaller batches give the writer room to checkpoint and vacuum between flushes, and ops can tune both without a deploy. We also reuse a single admin connection across all batches instead of reconnecting per batch — the tenant GUC is re-applied inside each batch's transaction sinceSET_CONFIG_QUERYisis_local=TRUE. And on a scan's first run we skip the pre-insertDELETEentirely (there's nothing to delete), so we don't open an empty write transaction or generate dead tuples for no reason.Partitioning.
compliance_requirements_overviewsis now a partitioned table, RANGE by its UUIDv7id, the same wayfindingsworks. That lets old scans' rows age out withDROP PARTITIONinstead ofDELETE+ CASCADE, which is where most of the autovacuum churn came from. The row generator emitsuuid7()ids now so new rows land in the right monthly partition (uuid4 would always fall intodefault). The old 5-column unique constraint is gone — once the partition keyidgets added to it, it doesn't mean anything, and per-scan dedup is already handled by the delete-then-reinsert, exactly likefindingsdoes.Migration
0098. Converts the table in place, without copying any data. It's aSeparateDatabaseAndState: the state side mirrors the model changes, and the database side renames the existing table and re-attaches it as thedefaultpartition of a new partitioned parent (ATTACHis metadata-only, no row rewrite). The parent FKs are added while it's still empty and the existing index is attached as-is, so nothing gets re-scanned or re-indexed. RLS is re-established on both the parent and the default partition through the project'sRowLevelSecurityConstraint.Steps to review
Please add a detailed description of how to review this PR.
Checklist
Community Checklist
SDK/CLI
UI
API
License
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
Summary by CodeRabbit
New Features
Bug Fixes