feat(api): backfill provider_str shadow column per-tenant

- Add batched per-tenant backfill job for the provider_str column
- Register backfill task on the backfill queue
- Dispatch the backfill from a data-only migration
This commit is contained in:
StylusFrost committed 2026-05-31 23:33:23 +02:00
1 parent 383e9c6bd8
commit 7dc0895581
4 files changed
+106

No files matched your search

@@ -0,0 +1,31 @@
from django.db import migrations
from tasks.tasks import backfill_provider_str_task
from api.db_router import MainRouter
def trigger_provider_str_backfill(apps, _schema_editor):
"""Dispatch a per-tenant Celery task to populate the transitional
`provider_str` shadow column for rows created before the sync trigger
from 0094 existed.
New writes are already covered by the trigger, so this only fills the gap
left by pre-existing rows. The work runs in the background, batched per
tenant, so the migration itself finishes in seconds regardless of table
size.
"""
Tenant = apps.get_model("api", "Tenant")
tenant_ids = Tenant.objects.using(MainRouter.admin_db).values_list("id", flat=True)
for tenant_id in tenant_ids:
backfill_provider_str_task.delay(tenant_id=str(tenant_id))
class Migration(migrations.Migration):
dependencies = [
("api", "0094_provider_str_shadow_column"),
]
operations = [
migrations.RunPython(trigger_provider_str_backfill, migrations.RunPython.noop),
]
+39
View File
@@ -543,3 +543,42 @@ class TestProviderShadowColumn:
provider.save()
provider.refresh_from_db()
assert provider.provider_str == provider.provider
def test_backfill_populates_pre_trigger_rows(self, providers_fixture):
"""Rows created before the sync trigger existed have a NULL shadow
column; the backfill copies `provider::text` into them in batches."""
from django.db import connection
from tasks.jobs.backfill import backfill_provider_str
tenant_id = str(providers_fixture[0].tenant_id)
# Simulate rows that predate the trigger: bypass triggers for the
# session, null the shadow column, then restore normal replication so
# subsequent writes resync as in production.
with connection.cursor() as cursor:
cursor.execute("SET session_replication_role = replica")
cursor.execute("UPDATE providers SET provider_str = NULL")
cursor.execute("SET session_replication_role = origin")
for provider in providers_fixture:
provider.refresh_from_db()
assert provider.provider_str is None
result = backfill_provider_str(tenant_id, batch_size=2)
assert result["updated"] == len(providers_fixture)
for provider in providers_fixture:
provider.refresh_from_db()
assert provider.provider_str == provider.provider
def test_backfill_is_idempotent(self, providers_fixture):
"""Re-running the backfill touches no rows once the column is populated
(the trigger already keeps it in sync), so it is safe to retry."""
from tasks.jobs.backfill import backfill_provider_str
tenant_id = str(providers_fixture[0].tenant_id)
result = backfill_provider_str(tenant_id)
assert result["updated"] == 0
+29
View File
@@ -40,6 +40,35 @@ from api.models import (
logger = get_task_logger(__name__)
def backfill_provider_str(tenant_id: str, batch_size: int = 1000):
"""Populate the transitional `provider_str` shadow column for rows that
predate the sync trigger, copying `provider::text` in bounded batches.
Each batch runs in its own RLS transaction so the lock is held only for the
rows in that batch, keeping the operation safe on a large table. Idempotent:
only rows where `provider_str IS NULL` are touched, so an interrupted run
resumes cleanly and a completed column is a no-op on retry.
"""
total_updated = 0
while True:
with rls_transaction(tenant_id) as cursor:
cursor.execute(
"UPDATE providers SET provider_str = provider::text "
"WHERE id IN ("
" SELECT id FROM providers WHERE provider_str IS NULL LIMIT %s"
")",
[batch_size],
)
updated = cursor.rowcount
total_updated += updated
if updated < batch_size:
break
logger.info(
"Backfilled provider_str for tenant %s: %d rows", tenant_id, total_updated
)
return {"tenant_id": tenant_id, "updated": total_updated}
def backfill_resource_scan_summaries(tenant_id: str, scan_id: str):
with rls_transaction(tenant_id, using=READ_REPLICA_ALIAS):
if ResourceScanSummary.objects.filter(
+7
View File
@@ -19,6 +19,7 @@ from tasks.jobs.backfill import (
backfill_daily_severity_summaries,
backfill_finding_group_summaries,
backfill_provider_compliance_scores,
backfill_provider_str,
backfill_resource_scan_summaries,
aggregate_scan_category_summaries,
aggregate_scan_resource_group_summaries,
@@ -722,6 +723,12 @@ def backfill_finding_group_summaries_task(tenant_id: str, days: int = None):
return backfill_finding_group_summaries(tenant_id=tenant_id, days=days)
@shared_task(name="backfill-provider-str", queue="backfill")
def backfill_provider_str_task(tenant_id: str):
"""Backfill the transitional provider_str shadow column for a tenant."""
return backfill_provider_str(tenant_id=tenant_id)
@shared_task(name="scan-category-summaries", queue="overview")
@handle_provider_deletion
def aggregate_scan_category_summaries_task(tenant_id: str, scan_id: str):