diff --git a/backend/kernelCI_app/management/commands/backfill_lab_ids.py b/backend/kernelCI_app/management/commands/backfill_lab_ids.py new file mode 100644 index 000000000..bc1c0c7ea --- /dev/null +++ b/backend/kernelCI_app/management/commands/backfill_lab_ids.py @@ -0,0 +1,161 @@ +import time +from collections.abc import Iterator + +from django.core.management.base import BaseCommand, CommandError +from django.db import connections, transaction +from django.db.backends.utils import CursorWrapper + +from kernelCI_app.constants.ingester import AUTOMATIC_LAB_FIELD, AUTOMATIC_LABS +from kernelCI_app.helpers.logger import out + + +def _real_lab_filter(misc_key: str) -> str: + """Keep non-automatic labs only (`_real_lab`). Named param: pattern.""" + return f""" + misc->>'{misc_key}' IS NOT NULL + AND misc->>'{misc_key}' !~ %(pattern)s + AND misc->>'{AUTOMATIC_LAB_FIELD}' IS NULL + """ + + +def _collect_lab_names(cur: CursorWrapper) -> list[str]: + cur.execute( + f""" + SELECT DISTINCT lab FROM ( + SELECT misc->>'lab' AS lab FROM builds + WHERE lab_id IS NULL + AND {_real_lab_filter("lab")} + UNION + SELECT misc->>'runtime' AS lab FROM tests + WHERE lab_id IS NULL + AND {_real_lab_filter("runtime")} + ) sub + """, + {"pattern": AUTOMATIC_LABS.pattern}, + ) + return [row[0] for row in cur.fetchall()] + + +def _ensure_labs(cur: CursorWrapper, lab_names: list[str]) -> None: + cur.executemany( + "INSERT INTO labs (name) VALUES (%s) ON CONFLICT (name) DO NOTHING", + [(name,) for name in lab_names], + ) + transaction.commit() + + +def _stage_pending_ids(cur: CursorWrapper, table: str, misc_key: str) -> str: + temp = f"_backfill_{table}_ids" + cur.execute(f"DROP TABLE IF EXISTS {temp}") + cur.execute( + f""" + CREATE TEMP TABLE {temp} AS + SELECT id FROM {table} + WHERE lab_id IS NULL + AND {_real_lab_filter(misc_key)} + """, + {"pattern": AUTOMATIC_LABS.pattern}, + ) + return temp + + +def _count_pending(cur: CursorWrapper, temp: str) -> int: + cur.execute(f"SELECT COUNT(*) FROM {temp}") + return cur.fetchone()[0] + + +def _backfill_table( + cur: CursorWrapper, + table: str, + misc_key: str, + temp: str, + batch_size: int, +) -> Iterator[int]: + """Yield cumulative updated row count after each batch.""" + total = 0 + while True: + cur.execute( + f""" + WITH batch AS ( + DELETE FROM {temp} + WHERE id IN (SELECT id FROM {temp} LIMIT %s) + RETURNING id + ), + updated_rows AS ( + UPDATE {table} AS target + SET lab_id = lab.id + FROM batch AS batch_row, labs AS lab + WHERE target.id = batch_row.id + AND target.misc->>%s = lab.name + RETURNING target.id + ) + SELECT + (SELECT COUNT(*) FROM batch), + (SELECT COUNT(*) FROM updated_rows) + """, + [batch_size, misc_key], + ) + batch_count, updated = cur.fetchone() + if batch_count == 0: + return + total += updated + transaction.commit() + yield total + + +class Command(BaseCommand): + help = ( + "Backfill lab_id FK on builds and tests from JSONB misc fields " + "(builds: misc->>'lab', tests: misc->>'runtime'). " + "Skips automatic labs (shell/k8s* or misc.automatic_lab); leaves lab_id NULL." + ) + + def add_arguments(self, parser) -> None: + parser.add_argument( + "--batch-size", + type=int, + default=50_000, + help="Rows per UPDATE batch (default: 50000)", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help=( + "Stage pending ids into a temp table and print counts, " + "but do not insert labs or update lab_id" + ), + ) + + def handle(self, *args, **options) -> None: + batch_size: int = options["batch_size"] + dry_run: bool = options["dry_run"] + + if batch_size < 1: + raise CommandError("--batch-size must be >= 1") + + with connections["default"].cursor() as cur: + if not dry_run: + out("Collecting distinct lab names from JSONB...") + lab_names = _collect_lab_names(cur) + out(f" Found {len(lab_names)} distinct lab names") + if lab_names: + out("Inserting missing lab names into labs table...") + _ensure_labs(cur, lab_names) + + for table, misc_key in (("builds", "lab"), ("tests", "runtime")): + out(f"Staging {table} ids...") + temp = _stage_pending_ids(cur, table, misc_key) + pending_count = _count_pending(cur, temp) + out(f"{table.capitalize()} to backfill: {pending_count}") + + if not pending_count or dry_run: + continue + + started_at = time.time() + for total in _backfill_table(cur, table, misc_key, temp, batch_size): + out( + f" {table}: {total}/{pending_count} " + f"({time.time() - started_at:.1f}s)" + ) + + out("Dry run complete, no lab_id changes made." if dry_run else "Done.")