feat: change ingester data insertion policy

Now allows for overwrites of null data

Closes #1552

Author: MarceloRobert

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

@@ -21,7 +21,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 +111,92 @@ def prepare_file_data(
         }
 
 
-def consume_buffer(buffer: list[TableModels], item_type: TableNames) -> None:
+def _generate_model_insert_query(
+    table_name: TableNames, model: TableModels
+) -> tuple[list[str], str]:
+    """
+    Dynamically generates the insert query for any model.
+    This function should only be used inside a transaction context.
+
+    Gives priority to the existing data in the database.
+
+    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:
+        conflict_clauses.append(
+            f"""
+            {field} = {op}({table_name}.{field}, EXCLUDED.{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(
+            "Unknown table '%s' passed to consume_buffer. Valid tables: %s"
+            % (table_name, str(", ".join(MODEL_MAP.keys())))
+        )
+        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

@@ -206,6 +206,61 @@ def test_prepare_file_data_file_error(self, mock_file_open, mock_logger):
         mock_file_open.assert_called_once()
 
 
+class TestGenerateInsertQuery:
+    def test_generate_model_insert_query_success(self):
+        """Test _generate_model_insert_query."""
+        from kernelCI_app.management.commands.helpers.kcidbng_ingester import (
+            _generate_model_insert_query,
+        )
+
+        mock_model = MagicMock()
+
+        mock_field_timestamp = MagicMock()
+        mock_field_timestamp.name = "field_timestamp"
+        mock_field_timestamp.db_column = "_timestamp"
+        mock_field_timestamp.generated = False
+        mock_field_timestamp.get_internal_type.return_value = "DateTimeField"
+
+        mock_field_comment = MagicMock()
+        mock_field_comment.name = "comment"
+        mock_field_comment.db_column = None
+        mock_field_comment.generated = False
+        mock_field_comment.get_internal_type.return_value = "CharField"
+
+        mock_field3 = MagicMock()
+        mock_field3.name = "checkout"
+        mock_field3.db_column = None
+        mock_field3.generated = False
+        mock_field3.get_internal_type.return_value = "ForeignKey"
+
+        mock_field4 = MagicMock()
+        mock_field4.name = "series"
+        mock_field4.generated = True
+
+        mock_model._meta.fields = [
+            mock_field_timestamp,
+            mock_field_comment,
+            mock_field3,
+            mock_field4,
+        ]
+
+        with patch(
+            "kernelCI_app.management.commands.helpers.kcidbng_ingester.MODEL_MAP",
+            {"issues": mock_model},
+        ):
+            updateable_fields, query = _generate_model_insert_query(
+                "issues", mock_model
+            )
+
+        assert updateable_fields == ["field_timestamp", "comment", "checkout_id"]
+        assert "INSERT INTO issues" in query
+        assert "_timestamp, comment, checkout_id" in query
+        assert "GREATEST(issues._timestamp, EXCLUDED._timestamp)" in query
+        assert "COALESCE(issues.comment, EXCLUDED.comment)" in query
+        assert "COALESCE(issues.checkout_id, EXCLUDED.checkout_id)" in query
+        assert "series" not in query
+
+
 class TestConsumeBuffer:
     """Test cases for consume_buffer function."""
 
@@ -218,24 +273,45 @@ class TestConsumeBuffer:
         INGEST_BATCH_SIZE_MOCK,
     )
     @patch("kernelCI_app.management.commands.helpers.kcidbng_ingester.out")
+    @patch(
+        "kernelCI_app.management.commands.helpers.kcidbng_ingester._generate_model_insert_query"
+    )
+    @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_generate_query, mock_out
+    ):
         """Test consume_buffer with items in buffer."""
+        table_name = "issues"
         mock_model = MagicMock()
         mock_buffer = [MagicMock(), MagicMock()]
+        mock_generate_query.return_value = (
+            ["_timestamp", "other_field"],
+            """
+                INSERT INTO issues (
+                    _timestamp, other_field
+                )
+                VALUES (
+                    %s, %s
+                )
+                ON CONFLICT (id)
+                DO UPDATE SET
+                    GREATEST(issues._timestamp, EXCLUDED._timestamp),
+                    COALESCE(issues.other_field, EXCLUDED.other_field);""",
+        )
+        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")
@@ -254,6 +330,12 @@ def test_consume_buffer_empty_buffer(self, mock_time, mock_out):
         mock_time.assert_not_called()
         mock_out.assert_not_called()
 
+    def test_consume_buffer_wrong_table(self):
+        """Test consume_buffer with invalid table name raises KeyError."""
+        with pytest.raises(KeyError):
+            mock_model = MagicMock()
+            consume_buffer([mock_model], "another")
+
 
 class TestFlushBuffers:
     """Test cases for flush_buffers function."""