Update "bulk_create_and_fetch" helper (#13068)

* Update "bulk_create_and_fetch" helper

Use bulk_create if DB backend supports returning

* Bug fix

src/backend/InvenTree/InvenTree/helpers_db.py

@@ -3,10 +3,19 @@
 import uuid
 from typing import Optional
 
-from django.db import transaction
+from django.db import connections, router, transaction
 from django.db.models import QuerySet
 
 
+def bulk_insert_returns_pks(model) -> bool:
+    """Return True if the database backend returns primary keys from bulk_create() for this model.
+
+    e.g. PostgreSQL, SQLite 3.35+ and MariaDB 10.5+ do, MySQL does not.
+    """
+    db = router.db_for_write(model)
+    return connections[db].features.can_return_rows_from_bulk_insert
+
+
 @transaction.atomic
 def bulk_create_and_fetch(
     model, items, id_field: str = 'pk', filters: Optional[dict] = None
@@ -20,34 +29,58 @@ def bulk_create_and_fetch(
         filters: Optional dictionary of filters to apply when fetching the created items.
 
     Returns:
-        A QuerySet containing the created items.
+        A QuerySet containing the created items, ordered by id_field (i.e. insertion order).
 
     This helper method is required because the Django bulk_create() method
     does not guarantee that the ID values of the created items will be populated in the returned objects.
     In particular, MySQL does not support returning the ID values of bulk created items.
 
-    So, we provide temporary metadata to the created items,
+    If the database backend returns primary keys from bulk_create() (e.g. PostgreSQL),
+    these are used directly. Otherwise, we provide temporary metadata to the created items,
     which can be used to fetch the created items from the database.
 
-    Assumptions:
+    Assumptions (for backends which do not return primary keys):
         - The provided model type has a "metadata" attribute which can be overloaded for this purpose
-        - No "metadata" is provided in the input items, as this will be overwritten by the method
         - The model type has an incrementing ID field (default: "pk")
 
+    Any "metadata" provided on the input items is retained (on both code paths).
+    """
+    id_field = id_field or 'pk'
+
+    if not items:
+        return model.objects.none()
+
+    if bulk_insert_returns_pks(model):
+        created = model.objects.bulk_create(items, batch_size=500)
+        pks = [item.pk for item in created]
+
+        return model.objects.filter(pk__in=pks, **(filters or {})).order_by(id_field)
+
+    return _bulk_create_and_fetch_by_metadata(model, items, id_field, filters)
+
+
+def _bulk_create_and_fetch_by_metadata(
+    model, items, id_field: str, filters: Optional[dict]
+) -> QuerySet:
+    """Fallback for bulk_create_and_fetch(), for backends which do not return primary keys.
+
+    The created items are tagged with temporary metadata, which is used to fetch them back.
     """
     bulk_create_id = uuid.uuid4().hex
 
     # Generate temporary metadata for bulk fetching
     metadata = {'bulk_create_id': bulk_create_id}
 
-    lookup_filters = dict(filters) if filters else {}
+    # Note: caller-provided filters are only applied to the returned queryset,
+    # so that *every* created item is found (and its temporary metadata cleared)
+    lookup_filters = {'metadata__bulk_create_id': bulk_create_id}
 
-    lookup_filters['metadata__bulk_create_id'] = bulk_create_id
+    # Find the "most recent" item in the database, to set a search floor
+    if instance := model.objects.order_by(f'-{id_field}').first():
+        lookup_filters[f'{id_field}__gt'] = getattr(instance, id_field)
 
-    if id_field:
-        # Find the "most recent" item in the database, to set a search floor
-        if instance := model.objects.order_by(f'-{id_field}').first():
-            lookup_filters[f'{id_field}__gt'] = getattr(instance, id_field)
+    # Retain any metadata provided with the input items, so that it can be restored
+    original_metadata = [item.metadata for item in items]
 
     # Overwrite the metadata values
     for item in items:
@@ -57,16 +90,31 @@ def bulk_create_and_fetch(
 
     instances = model.objects.filter(**lookup_filters)
 
-    id_field = id_field or 'pk'
-
     pks = list(instances.order_by(id_field).values_list(id_field, flat=True))
 
     # Override the metadata values to remove the temporary bulk_create_id
     instances.update(metadata=None)
 
+    # Restore the original metadata on the input items
+    for item, original in zip(items, original_metadata, strict=True):
+        item.metadata = original
+
+    # ...and in the database, for any items which provided their own metadata.
+    # Note: pks are in insertion order, so correspond positionally with the input items
+    restored = [
+        model(pk=pk, metadata=original)
+        for pk, original in zip(pks, original_metadata, strict=True)
+        if original is not None
+    ]
+
+    if restored:
+        model.objects.bulk_update(restored, ['metadata'], batch_size=500)
+
     # Fetch the newly created items (by primary key, as the metadata filter no longer matches)
     # Ordered by id_field: auto-increment values are assigned in insertion order on every
     # backend, so this restores correspondence with the input 'items' list for callers that
     # zip the result against it positionally. Without this, some backends (e.g. MySQL) may
     # return 'pk__in' rows in a different order than they were inserted.
-    return model.objects.filter(**{f'{id_field}__in': pks}).order_by(id_field)
+    return model.objects.filter(**{f'{id_field}__in': pks}, **(filters or {})).order_by(
+        id_field
+    )

src/backend/InvenTree/InvenTree/test_helpers_db.py

@@ -1,11 +1,13 @@
 """Functional tests for the InvenTree.helpers_db module."""
 
+from unittest import mock
+
 from django.db import connection
 from django.test import TestCase
 from django.test.utils import CaptureQueriesContext
 
 from company.models import Company, Contact
-from InvenTree.helpers_db import bulk_create_and_fetch
+from InvenTree.helpers_db import bulk_create_and_fetch, bulk_insert_returns_pks
 from order.models import PurchaseOrder
 from part.models import Part
 from stock.models import StockItem
@@ -74,6 +76,40 @@ class BulkCreateAndFetchTest(TestCase):
         self.assertEqual(len(result), 2)
         self.assertNotIn('Existing', [c.name for c in result])
 
+    def test_preserves_metadata(self):
+        """Metadata provided with the input items is retained (in the database and on the items)."""
+        company = Company.objects.create(name='ACME')
+
+        items = [
+            Contact(company=company, name='With metadata', metadata={'key': 'value'}),
+            Contact(company=company, name='Without metadata'),
+        ]
+
+        result = bulk_create_and_fetch(Contact, items)
+
+        self.assertEqual(result[0].metadata, {'key': 'value'})
+        self.assertIsNone(result[1].metadata)
+
+        # The input items are not left with the temporary bulk-create metadata
+        self.assertEqual(items[0].metadata, {'key': 'value'})
+        self.assertIsNone(items[1].metadata)
+
+        # No temporary metadata remains in the database
+        self.assertFalse(
+            Contact.objects.filter(metadata__has_key='bulk_create_id').exists()
+        )
+
+    def test_insertion_order(self):
+        """Results are returned in insertion order (callers zip them against the input)."""
+        company = Company.objects.create(name='ACME')
+
+        names = [f'Contact {idx:03d}' for idx in reversed(range(50))]
+        items = [Contact(company=company, name=name) for name in names]
+
+        result = bulk_create_and_fetch(Contact, items)
+
+        self.assertEqual([c.name for c in result], names)
+
     def test_empty_list(self):
         """Calling with an empty list of items should not raise, and return no items."""
         result = bulk_create_and_fetch(Contact, [])
@@ -161,3 +197,67 @@ class BulkCreateAndFetchTest(TestCase):
 
         # Perform a very large bulk-create - this will be batched
         self._create_and_count_queries(2000)
+
+
+class BulkCreateAndFetchFallbackTest(BulkCreateAndFetchTest):
+    """Run all bulk_create_and_fetch tests against the fallback (metadata lookup) code path.
+
+    This path is used for database backends which do not return primary keys from
+    bulk_create() (e.g. MySQL) - forcing it here ensures it is tested on every backend.
+    """
+
+    def setUp(self):
+        """Force the fallback code path."""
+        super().setUp()
+
+        patcher = mock.patch(
+            'InvenTree.helpers_db.bulk_insert_returns_pks', return_value=False
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
+
+
+class BulkCreateAndFetchPathTest(TestCase):
+    """Tests for the selection of the bulk_create_and_fetch code path."""
+
+    def test_backend_detection(self):
+        """The fast path is selected according to the database backend features."""
+        self.assertEqual(
+            bulk_insert_returns_pks(Contact),
+            connection.features.can_return_rows_from_bulk_insert,
+        )
+
+    def test_fast_path(self):
+        """When primary keys are returned, the temporary metadata is not used at all."""
+        if not bulk_insert_returns_pks(Contact):  # pragma: no cover
+            self.skipTest('Database backend does not return primary keys')
+
+        company = Company.objects.create(name='ACME')
+        items = [Contact(company=company, name=f'Contact {i}') for i in range(5)]
+
+        with CaptureQueriesContext(connection) as ctx:
+            result = list(bulk_create_and_fetch(Contact, items))
+
+        self.assertEqual(len(result), 5)
+
+        # Ignore transaction management statements (from @transaction.atomic)
+        queries = [
+            q['sql']
+            for q in ctx.captured_queries
+            if 'SAVEPOINT' not in q['sql'].upper()
+        ]
+
+        # Insert + fetch only - no metadata tagging / clearing
+        self.assertEqual(len(queries), 2, queries)
+        self.assertFalse(any('bulk_create_id' in sql for sql in queries))
+
+    def test_fast_path_filters(self):
+        """Caller-provided filters are applied on the fast path too."""
+        company_a = Company.objects.create(name='A')
+        company_b = Company.objects.create(name='B')
+
+        items = [Contact(company=company_a, name='Contact')]
+
+        result = bulk_create_and_fetch(Contact, items, filters={'company': company_b})
+
+        self.assertEqual(len(result), 0)