diff --git a/src/backend/InvenTree/InvenTree/helpers_db.py b/src/backend/InvenTree/InvenTree/helpers_db.py index 5ed25e0bee..3388c5fb04 100644 --- a/src/backend/InvenTree/InvenTree/helpers_db.py +++ b/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 + ) diff --git a/src/backend/InvenTree/InvenTree/test_helpers_db.py b/src/backend/InvenTree/InvenTree/test_helpers_db.py index de4da41b50..3d7b40ba2c 100644 --- a/src/backend/InvenTree/InvenTree/test_helpers_db.py +++ b/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)