From 338b2cd79e362a09c4d47f5adb5bdff77668d403 Mon Sep 17 00:00:00 2001 From: Alan Peixinho Date: Fri, 10 Jul 2026 18:16:45 -0300 Subject: [PATCH 1/2] feat: Add backfill labs command * process historical data to use the lab foreign key Signed-off-by: Alan Peixinho --- .../management/commands/backfill_lab_ids.py | 149 ++++++++++++++++++ 1 file changed, 149 insertions(+) create mode 100644 backend/kernelCI_app/management/commands/backfill_lab_ids.py 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..64c134097 --- /dev/null +++ b/backend/kernelCI_app/management/commands/backfill_lab_ids.py @@ -0,0 +1,149 @@ +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.helpers.logger import out + + +def _collect_lab_names(cur: CursorWrapper) -> list[str]: + cur.execute( + """ + SELECT DISTINCT lab FROM ( + SELECT misc->>'lab' AS lab FROM builds + WHERE lab_id IS NULL + AND misc->>'lab' IS NOT NULL + UNION + SELECT misc->>'runtime' AS lab FROM tests + WHERE lab_id IS NULL + AND misc->>'runtime' IS NOT NULL + ) sub + """ + ) + 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 misc->>%s IS NOT NULL + """, + [misc_key], + ) + 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')." + ) + + 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.") From 28acd8b8c0b7efe63f1e4601016bdd1984be834c Mon Sep 17 00:00:00 2001 From: Alan Peixinho Date: Mon, 13 Jul 2026 18:15:48 -0300 Subject: [PATCH 2/2] fix: skip automatic labs in lab_id backfill Leave lab_id NULL for shell/k8s* and misc.automatic_lab so backfill matches ingest _real_lab and does not insert fake labs. Signed-off-by: Alan Peixinho --- .../management/commands/backfill_lab_ids.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/backend/kernelCI_app/management/commands/backfill_lab_ids.py b/backend/kernelCI_app/management/commands/backfill_lab_ids.py index 64c134097..bc1c0c7ea 100644 --- a/backend/kernelCI_app/management/commands/backfill_lab_ids.py +++ b/backend/kernelCI_app/management/commands/backfill_lab_ids.py @@ -5,22 +5,33 @@ 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 misc->>'lab' IS NOT NULL + AND {_real_lab_filter("lab")} UNION SELECT misc->>'runtime' AS lab FROM tests WHERE lab_id IS NULL - AND misc->>'runtime' IS NOT NULL + AND {_real_lab_filter("runtime")} ) sub - """ + """, + {"pattern": AUTOMATIC_LABS.pattern}, ) return [row[0] for row in cur.fetchall()] @@ -41,9 +52,9 @@ def _stage_pending_ids(cur: CursorWrapper, table: str, misc_key: str) -> str: CREATE TEMP TABLE {temp} AS SELECT id FROM {table} WHERE lab_id IS NULL - AND misc->>%s IS NOT NULL + AND {_real_lab_filter(misc_key)} """, - [misc_key], + {"pattern": AUTOMATIC_LABS.pattern}, ) return temp @@ -95,7 +106,8 @@ def _backfill_table( class Command(BaseCommand): help = ( "Backfill lab_id FK on builds and tests from JSONB misc fields " - "(builds: misc->>'lab', tests: misc->>'runtime')." + "(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: