diff --git a/dojo/location/models.py b/dojo/location/models.py index 99fc308886..28dfd8a28f 100644 --- a/dojo/location/models.py +++ b/dojo/location/models.py @@ -6,7 +6,7 @@ from django.contrib.postgres.indexes import GinIndex from django.contrib.postgres.search import SearchVector from django.core.validators import MinLengthValidator -from django.db import transaction +from django.db import IntegrityError, transaction from django.db.models import ( CASCADE, RESTRICT, @@ -410,6 +410,20 @@ def get_or_create_from_object(cls, location: Self) -> Self: msg = "Subclasses must implement get_or_create_from_object" raise NotImplementedError(msg) + # Upper bound on the concurrent-writer recovery retries in + # bulk_get_or_create. Each pass shrinks the batch to only the rows still + # missing, so the loop makes progress every iteration; the cap guards + # against pathological churn rather than being expected to be reached. + _BULK_CREATE_MAX_ATTEMPTS = 3 + + @classmethod + def _existing_by_identity_hash(cls, hashes: list[str]) -> dict[str, Self]: + """Map identity_hash -> already-saved instance for the given hashes.""" + return { + obj.identity_hash: obj + for obj in cls.objects.filter(identity_hash__in=hashes).select_related("location") + } + @classmethod def bulk_get_or_create(cls, locations: Iterable[Self]) -> list[Self]: """ @@ -419,6 +433,11 @@ def bulk_get_or_create(cls, locations: Iterable[Self]) -> list[Self]: bulk_create for both the parent Location rows and the subtype rows. Returns the full list of saved instances (existing + newly created), in the same order as the input. Duplicate inputs map to the same saved instance. + + identity_hash is a global singleton, so two imports that reference the same + package (Dependency) or endpoint (URL) race to create the same row. Creation + tolerates that race — see ``_bulk_create_tolerating_races`` — rather than + letting one importer abort the other's whole scan import. """ if not locations: return [] @@ -435,52 +454,107 @@ def bulk_get_or_create(cls, locations: Iterable[Self]) -> list[Self]: hashes.append(loc.identity_hash) # Look up existing objects, grouping by hash - existing_by_hash = { - obj.identity_hash: obj - for obj in cls.objects.filter(identity_hash__in=hashes).select_related("location") - } + existing_by_hash = cls._existing_by_identity_hash(hashes) - # Create the list of new locations to create - new_locations = [] + # Determine which locations still need creating, deduplicated by hash. + to_create: dict[str, Self] = {} for loc in locations: - if loc.identity_hash not in existing_by_hash: - new_locations.append(loc) - # Mark it so we don't try to create duplicates within the same batch - existing_by_hash[loc.identity_hash] = loc - else: + if loc.identity_hash in existing_by_hash: # Preserve association data from the input onto the existing saved object, in case we're associating # existing locations with findings/products saved = existing_by_hash[loc.identity_hash] if hasattr(loc, "_association_data") and not hasattr(saved, "_association_data"): saved._association_data = loc._association_data + elif loc.identity_hash not in to_create: + # First occurrence of a not-yet-persisted hash in this batch + to_create[loc.identity_hash] = loc - # Create 'em - if new_locations: - location_type = cls.get_location_type() - with transaction.atomic(): - # Bulk create parent Locations - parents = [ - Location( - location_type=location_type, - location_value=loc.get_location_value(), - ) - for loc in new_locations - ] - Location.objects.bulk_create(parents, batch_size=1000) - # Assign Location FKs to the subtypes, then bulk create them. - for loc, parent in zip(new_locations, parents, strict=True): - loc.location_id = parent.id - loc.location = parent - # Note: there is a subtle potential race condition here, if somehow one of the locations to be created - # has already been created, e.g. by a separate thread that commits while this thread is running. Setting - # `ignore_conflicts=True` here would prevent this step from raising an IntegrityError, but would leave - # dangling parent Location objects that were created above. Rather than performing a cleanup in that - # (unlikely?) case, just allow the transaction to rollback. - cls.objects.bulk_create(new_locations, batch_size=1000) + # Create 'em (tolerating a concurrent writer that beats us to some rows) + if to_create: + existing_by_hash.update(cls._bulk_create_tolerating_races(to_create)) # Return in input order return [existing_by_hash[h] for h in hashes] + @classmethod + def _bulk_create_tolerating_races(cls, to_create: dict[str, Self]) -> dict[str, Self]: + """ + Bulk create the given subtype rows (keyed by identity_hash), tolerating a + concurrent writer that commits some of the same identity_hashes first. + + Parent Location rows and their subtype rows are created together inside a + savepoint. A unique-constraint collision aborts the whole INSERT, so the + savepoint rolls back — the batch's parent Locations included, leaving no + orphaned rows — and the enclosing transaction stays usable. We then + re-resolve the rows the concurrent writer committed, drop them from the + batch, and retry only the genuine remainder. + + Without this, one racing import raised IntegrityError out of the whole + persist() and aborted the entire scan import: identity_hash is a global + singleton, so two imports referencing the same package collide routinely. + """ + location_type = cls.get_location_type() + resolved: dict[str, Self] = {} + remaining = dict(to_create) + + for _attempt in range(cls._BULK_CREATE_MAX_ATTEMPTS): + batch = list(remaining.values()) + try: + with transaction.atomic(): + # Bulk create parent Locations + parents = [ + Location( + location_type=location_type, + location_value=loc.get_location_value(), + ) + for loc in batch + ] + Location.objects.bulk_create(parents, batch_size=1000) + # Assign Location FKs to the subtypes, then bulk create them. + for loc, parent in zip(batch, parents, strict=True): + loc.location_id = parent.id + loc.location = parent + cls.objects.bulk_create(batch, batch_size=1000) + except IntegrityError: + # A concurrent writer committed one or more of these between our + # existence check and this INSERT. Re-resolve those, carry over + # their association data, and retry the rest. The savepoint + # rollback already discarded this batch's parents, so nothing is + # orphaned. + remaining = cls._absorb_raced_rows(remaining, resolved) + if not remaining: + return resolved + continue + # No collision: every remaining row was created in this batch. + resolved.update(remaining) + return resolved + + # Retries exhausted (persistent contention). Resolve whatever now exists + # so callers still get saved instances; only a hash that is genuinely + # still absent is a real failure. + remaining = cls._absorb_raced_rows(remaining, resolved) + if remaining: + error_message = ( + f"Could not persist {cls.__name__} rows for identity_hashes: {sorted(remaining)}" + ) + raise IntegrityError(error_message) + return resolved + + @classmethod + def _absorb_raced_rows(cls, remaining: dict[str, Self], resolved: dict[str, Self]) -> dict[str, Self]: + """ + Move rows a concurrent writer already committed out of ``remaining`` and + into ``resolved`` (carrying the input's association data onto the fetched + instance), and return the hashes still needing creation. + """ + now_existing = cls._existing_by_identity_hash(list(remaining)) + for identity_hash, saved in now_existing.items(): + source = remaining[identity_hash] + if hasattr(source, "_association_data") and not hasattr(saved, "_association_data"): + saved._association_data = source._association_data + resolved.update(now_existing) + return {h: loc for h, loc in remaining.items() if h not in now_existing} + class ReferenceDataMixin(Model): diff --git a/unittests/test_bulk_locations.py b/unittests/test_bulk_locations.py index 2917aa2fe5..e0a6a6ac80 100644 --- a/unittests/test_bulk_locations.py +++ b/unittests/test_bulk_locations.py @@ -139,6 +139,55 @@ def test_transaction_atomicity(self): self.assertEqual(Location.objects.count(), initial_count) + def test_recovers_when_row_committed_concurrently(self): + """ + A concurrent writer that commits the same identity_hash between our + existence check and our INSERT must not abort the batch. + + Reproduces the production failure (``duplicate key value violates unique + constraint "..._identity_hash_key"`` raised out of an entire scan import): + the pre-existing row is hidden from the FIRST existence lookup only, so + the code attempts to INSERT a duplicate and hits the real DB unique + constraint, then must re-resolve the existing row and still create the + genuinely-new one instead of raising. + """ + existing = URL.get_or_create_from_object(_make_url("oss-race-existing.example.com")) + incoming = [ + _make_url("oss-race-existing.example.com"), # committed by a "concurrent" writer + _make_url("oss-race-new.example.com"), # genuinely new + ] + + original_lookup = URL._existing_by_identity_hash + state = {"calls": 0} + + def racing_lookup(hashes): + state["calls"] += 1 + result = original_lookup(hashes) + if state["calls"] == 1: + # Simulate the row not yet being visible when we first checked + result.pop(existing.identity_hash, None) + return result + + with patch.object(URL, "_existing_by_identity_hash", side_effect=racing_lookup): + saved = URL.bulk_get_or_create(incoming) + + by_hash = {s.identity_hash: s for s in saved} + # Order preserved and every input resolved to a saved instance + self.assertEqual(len(saved), 2) + self.assertTrue(all(s.pk is not None and s.location_id is not None for s in saved)) + # The raced row resolves to the ORIGINAL, with no duplicate created + self.assertEqual(by_hash[existing.identity_hash].pk, existing.pk) + self.assertEqual(URL.objects.filter(identity_hash=existing.identity_hash).count(), 1) + # The genuinely-new row was still created + self.assertNotEqual(by_hash[incoming[1].identity_hash].pk, existing.pk) + # No orphaned parent Location rows left behind by the rolled-back batch + orphan_urls = ( + Location.objects.filter(location_type="url") + .exclude(pk__in=URL.objects.values("location_id")) + .exists() + ) + self.assertFalse(orphan_urls) + # --------------------------------------------------------------------------- # LocationManager._bulk_get_or_create_locations (URL-only)