feat: change ingester data insertion policy

Now allows for overwrites of null data

Closes #1552

Author: MarceloRobert

backend/kernelCI_app/constants/ingester.py

@@ -48,3 +48,9 @@
 except (ValueError, TypeError):
     logger.warning("Invalid INGEST_QUEUE_MAXSIZE, using default 5000")
     INGEST_QUEUE_MAXSIZE = 5000
+
+
+PRIO_DB = is_boolean_or_string_true(os.environ.get("PRIO_BD", "True"))
+"""Toggles the priority when updating database data.\n
+If True, this will prioritize what is already present in the database, never updating fields != null.
+If False, it will prioritize the incoming data, allowing it to overwrite existing data."""

backend/kernelCI_app/management/commands/helpers/kcidbng_ingester.py

@@ -10,6 +10,7 @@
     INGEST_FLUSH_TIMEOUT_SEC,
     INGEST_QUEUE_MAXSIZE,
     VERBOSE,
+    PRIO_DB,
 )
 import threading
 import time
@@ -21,7 +22,7 @@
     extract_log_excerpt,
 )
 import kcidb_io
-from django.db import transaction
+from django.db import connections, transaction
 from kernelCI_app.models import Issues, Checkouts, Builds, Tests, Incidents
 
 from kernelCI_app.management.commands.helpers.process_submissions import (
@@ -111,23 +112,95 @@ def prepare_file_data(
         }
 
 
-def consume_buffer(buffer: list[TableModels], item_type: TableNames) -> None:
+def _generate_model_insert_query(
+    table_name: TableModels, model: TableModels
+) -> tuple[list[str], str]:
+    """
+    Dynamically generates the insert query for any model.
+
+    The generated query follows the policy of never updating non-null values.
+    This policy can be changed with the `PRIO_DB` env var.
+
+    Returns a list of which model properties can be updated and the insert query.
+    """
+    updateable_model_fields: list[str] = []
+    updateable_db_fields: list[str] = []
+    query_params_properties: list[tuple[str, str]] = []
+
+    for field in model._meta.fields:
+        if field.generated:
+            continue
+
+        field_name = (
+            field.name + "_id"
+            if field.get_internal_type() == "ForeignKey"
+            else field.name
+        )
+        real_name = field.db_column or field_name
+        operation = "GREATEST" if real_name == "_timestamp" else "COALESCE"
+
+        query_params_properties.append((real_name, operation))
+        updateable_model_fields.append(field_name)
+        updateable_db_fields.append(real_name)
+
+    conflict_clauses = []
+    for field, op in query_params_properties:
+        if PRIO_DB:
+            conflict_clauses.append(
+                f"""
+            {field} = {op}({table_name}.{field}, EXCLUDED.{field})"""
+            )
+        else:
+            conflict_clauses.append(
+                f"""
+            {field} = {op}(EXCLUDED.{field}, {table_name}.{field})"""
+            )
+
+    query = f"""
+        INSERT INTO {table_name} (
+            {', '.join(updateable_db_fields)}
+        )
+        VALUES (
+            {', '.join(['%s'] * len(updateable_db_fields))}
+        )
+        ON CONFLICT (id)
+        DO UPDATE SET {', '.join(conflict_clauses)};
+    """
+
+    return updateable_model_fields, query
+
+
+def consume_buffer(buffer: list[TableModels], table_name: TableNames) -> None:
     """
     Consume a buffer of items and insert them into the database.
     This function is called by the db_worker thread.
     """
     if not buffer:
         return
 
-    model = MODEL_MAP[item_type]
+    try:
+        model = MODEL_MAP[table_name]
+    except KeyError:
+        out(f"Unknown table '{table_name}' passed to consume_buffer")
+        raise
+
+    updateable_model_fields, query = _generate_model_insert_query(table_name, model)
+
+    params = []
+    for obj in buffer:
+        obj_values = []
+        for field in updateable_model_fields:
+            value = getattr(obj, field)
+            if isinstance(value, (dict, list)):
+                value = json.dumps(value)
+            obj_values.append(value)
+        params.append(tuple(obj_values))
 
     t0 = time.time()
-    model.objects.bulk_create(
-        buffer,
-        batch_size=INGEST_BATCH_SIZE,
-        ignore_conflicts=True,
-    )
-    out("bulk_create %s: n=%d in %.3fs" % (item_type, len(buffer), time.time() - t0))
+    with connections["default"].cursor() as cursor:
+        cursor.executemany(query, params)
+
+    out("bulk_create %s: n=%d in %.3fs" % (table_name, len(buffer), time.time() - t0))
 
 
 def flush_buffers(

backend/kernelCI_app/models.py

@@ -1,3 +1,6 @@
+"""Defines the models used in the main database.
+All models should have explicit id column for the ingester to work properly."""
+
 from django.db import models
 from django.contrib.postgres.fields import ArrayField
 from django.contrib.postgres.indexes import GinIndex

backend/kernelCI_app/tests/unitTests/commands/monitorSubmissions/kcidbng_ingester_test.py

@@ -218,24 +218,26 @@ class TestConsumeBuffer:
         INGEST_BATCH_SIZE_MOCK,
     )
     @patch("kernelCI_app.management.commands.helpers.kcidbng_ingester.out")
+    @patch("kernelCI_app.management.commands.helpers.kcidbng_ingester.connections")
     @patch("time.time", side_effect=TIME_MOCK)
-    def test_consume_buffer_with_items(self, mock_time, mock_out):
+    def test_consume_buffer_with_items(self, mock_time, mock_connections, mock_out):
         """Test consume_buffer with items in buffer."""
+        table_name = "issues"
         mock_model = MagicMock()
         mock_buffer = [MagicMock(), MagicMock()]
+        mock_cursor = MagicMock()
+        mock_connections["default"].cursor.return_value.__enter__.return_value = (
+            mock_cursor
+        )
 
         with patch(
             "kernelCI_app.management.commands.helpers.kcidbng_ingester.MODEL_MAP",
             {"issues": mock_model},
         ):
-            consume_buffer(mock_buffer, "issues")
+            consume_buffer(mock_buffer, table_name)
 
         assert mock_time.call_count == 2
-        mock_model.objects.bulk_create.assert_called_once_with(
-            mock_buffer,
-            batch_size=INGEST_BATCH_SIZE_MOCK,
-            ignore_conflicts=True,
-        )
+        mock_cursor.executemany.assert_called_once()
         mock_out.assert_called_once()
 
     @patch("kernelCI_app.management.commands.helpers.kcidbng_ingester.out")