feat(backup): Add ForeignKeyComparator (#54558)

Given two instances of a model that are being validated for equality, we
want the foreign keys to be correct relatively (ie, they point to the
same model in the respective models' JSON blobs), not absolutely (ie,
they are literally the same integer). By creating maps that store the
relations between pks and ordinals, we can easily check that the models
point to match their respective ordinals regardless of the actual pk
numbers.

Issue: getsentry/team-ospo#171

Author: azaslavsky

src/sentry/backup/comparators.py

@@ -7,6 +7,7 @@
 from dateutil import parser
 from django.db import models
 
+from sentry.backup.dependencies import PrimaryKeyMap, dependencies
 from sentry.backup.findings import ComparatorFinding, ComparatorFindingKind, InstanceID
 from sentry.backup.helpers import Side, get_exportable_final_derivations_of
 from sentry.db.models import BaseModel
@@ -200,6 +201,73 @@ def compare(self, on: InstanceID, left: JSONData, right: JSONData) -> list[Compa
         return findings
 
 
+class ForeignKeyComparator(JSONScrubbingComparator):
+    """Ensures that foreign keys match in a relative (they refer to the same other model in their
+    respective JSON blobs) rather than absolute (they have literally the same integer value)
+    sense."""
+
+    left_pk_map: PrimaryKeyMap | None = None
+    right_pk_map: PrimaryKeyMap | None = None
+
+    def __init__(self, foreign_fields: dict[str, models.base.ModelBase]):
+        super().__init__(*(foreign_fields.keys()))
+        self.foreign_fields = foreign_fields
+
+    def set_primary_key_maps(self, left_pk_map: PrimaryKeyMap, right_pk_map: PrimaryKeyMap):
+        """Call this function before running the comparator, to ensure that it has access to the latest mapping information for both sides of the comparison."""
+
+        self.left_pk_map = left_pk_map
+        self.right_pk_map = right_pk_map
+
+    def compare(self, on: InstanceID, left: JSONData, right: JSONData) -> list[ComparatorFinding]:
+        findings = []
+        fields = sorted(self.fields)
+        for f in fields:
+            field_model_name = "sentry." + self.foreign_fields[f].__name__.lower()
+            if left["fields"].get(f) is None and right["fields"].get(f) is None:
+                continue
+
+            if self.left_pk_map is None or self.right_pk_map is None:
+                raise RuntimeError("must call `set_primary_key_maps` before comparing")
+
+            left_fk_as_ordinal = self.left_pk_map.get(field_model_name, left["fields"][f])
+            right_fk_as_ordinal = self.right_pk_map.get(field_model_name, right["fields"][f])
+            if left_fk_as_ordinal is None or right_fk_as_ordinal is None:
+                if left_fk_as_ordinal is None:
+                    findings.append(
+                        ComparatorFinding(
+                            kind=self.get_kind(),
+                            on=on,
+                            left_pk=left["pk"],
+                            right_pk=right["pk"],
+                            reason=f"""the left foreign key ordinal for `{f}` model with pk `{left["fields"][f]}` could not be found""",
+                        )
+                    )
+                if right_fk_as_ordinal is None:
+                    findings.append(
+                        ComparatorFinding(
+                            kind=self.get_kind(),
+                            on=on,
+                            left_pk=left["pk"],
+                            right_pk=right["pk"],
+                            reason=f"""the right foreign key ordinal for `{f}` model with pk `{right["fields"][f]}` could not be found""",
+                        )
+                    )
+                continue
+
+            if left_fk_as_ordinal != right_fk_as_ordinal:
+                findings.append(
+                    ComparatorFinding(
+                        kind=self.get_kind(),
+                        on=on,
+                        left_pk=left["pk"],
+                        right_pk=right["pk"],
+                        reason=f"""the left foreign key ordinal ({left_fk_as_ordinal}) for `{f}` was not equal to the right foreign key ordinal ({right_fk_as_ordinal})""",
+                    )
+                )
+        return findings
+
+
 class ObfuscatingComparator(JSONScrubbingComparator, ABC):
     """Comparator that compares private values, but then safely truncates them to ensure that they
     do not leak out in logs, stack traces, etc."""
@@ -338,6 +406,16 @@ def auto_assign_email_obfuscating_comparators(comps: ComparatorMap) -> None:
                 comps[name].append(EmailObfuscatingComparator(*assign))
 
 
+def auto_assign_foreign_key_comparators(comps: ComparatorMap) -> None:
+    """Automatically assigns the ForeignKeyComparator or to all appropriate model fields (see
+    dependencies.py for more on what "appropriate" means in this context)."""
+
+    for model_name, rels in dependencies().items():
+        comps[model_name.lower()].append(
+            ForeignKeyComparator({k: v.model for k, v in rels.foreign_keys.items()})
+        )
+
+
 ComparatorList = List[JSONScrubbingComparator]
 ComparatorMap = Dict[str, ComparatorList]
 
@@ -380,6 +458,7 @@ def build_default_comparators():
     # to the `DEFAULT_COMPARATORS` map.
     auto_assign_datetime_equality_comparators(comparators)
     auto_assign_email_obfuscating_comparators(comparators)
+    auto_assign_foreign_key_comparators(comparators)
 
     return comparators
 

src/sentry/backup/dependencies.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from collections import defaultdict
 from enum import Enum, auto, unique
 from typing import NamedTuple
 
@@ -71,6 +72,37 @@ def default(self, obj):
         return super().default(obj)
 
 
+class PrimaryKeyMap:
+    """
+    A map between a primary key in one primary key sequence (like a database) and another (like the
+    ordinals in a backup JSON file). As models are moved between databases, their absolute contents
+    may stay the same, but their relative identifiers may change. This class allows us to track how
+    those relations have been transformed, thereby preserving the model references between one
+    another.
+
+    Note that the map assumes that the primary keys in question are integers. In particular, natural
+    keys are not supported!
+    """
+
+    mapping: dict[str, dict[int, int]]
+
+    def __init__(self):
+        self.mapping = defaultdict(dict)
+
+    def get(self, model: str, old: int) -> int | None:
+        """Get the new, post-mapping primary key from an old primary key."""
+
+        pk_map = self.mapping.get(model)
+        if pk_map is None:
+            return None
+        return pk_map.get(old)
+
+    def insert(self, model: str, old: int, new: int):
+        """Create a new OLD_PK -> NEW_PK mapping for the given model."""
+
+        self.mapping[model][old] = new
+
+
 def dependencies() -> dict[str, ModelRelations]:
     """Produce a dictionary mapping model type definitions to a `ModelDeps` describing their dependencies."""
 

src/sentry/backup/findings.py

@@ -62,6 +62,13 @@ class ComparatorFindingKind(IntEnum):
     # `None`.
     HashObfuscatingComparatorExistenceCheck = auto()
 
+    # Foreign key field comparison failed.
+    ForeignKeyComparator = auto()
+
+    # Failed to compare foreign key fields because one of the fields being compared was not present
+    # or `None`.
+    ForeignKeyComparatorExistenceCheck = auto()
+
     # Failed to compare an ignored field.
     IgnoredComparator = auto()
 

src/sentry/backup/validate.py

@@ -5,7 +5,8 @@
 from difflib import unified_diff
 from typing import Dict, Tuple
 
-from sentry.backup.comparators import DEFAULT_COMPARATORS, ComparatorMap
+from sentry.backup.comparators import DEFAULT_COMPARATORS, ComparatorMap, ForeignKeyComparator
+from sentry.backup.dependencies import PrimaryKeyMap
 from sentry.backup.findings import (
     ComparatorFinding,
     ComparatorFindingKind,
@@ -115,10 +116,30 @@ def json_lines(obj: JSONData) -> list[str]:
     if not findings.empty():
         return findings
 
+    # As models are compared, we will add their pk mapping to separate `PrimaryKeyMaps`. Then, when
+    # a foreign keyed field into the specific model is encountered, we will be able to ensure that
+    # both sides reference the correct model.
+    #
+    # For instance, we encounter the first `sentry.User` model on both the left and right side, with
+    # the left side having a `pk` of 123, and the right having `456`. This means that we want to map
+    # `[sentry.User][123] = 1` on the left and `[sentry.User][456] = 1`. Later, when we encounter
+    # foreign keys to a user model with `pk` 123 on the left and 456 on the right, we'll be able to
+    # dereference the map to ensure that those both point to the same model on their respective
+    # sides.
+    left_pk_map = PrimaryKeyMap()
+    right_pk_map = PrimaryKeyMap()
+
     # We only perform custom comparisons and JSON diffs on non-duplicate entries that exist in both
     # outputs.
     for id, right in right_models.items():
+        if id.ordinal is None:
+            raise RuntimeError("all InstanceIDs used for comparisons must have their ordinal set")
+
+        # Save the pk -> ordinal mapping on both sides, so that we can decode foreign keys into this
+        # model that we encounter later.
         left = left_models[id]
+        left_pk_map.insert(id.model, left["pk"], id.ordinal)
+        right_pk_map.insert(id.model, right["pk"], id.ordinal)
 
         # Try comparators applicable for this specific model.
         if id.model in comparators:
@@ -134,6 +155,9 @@ def json_lines(obj: JSONData) -> list[str]:
                     findings.extend(ex)
                     continue
 
+                if isinstance(cmp, ForeignKeyComparator):
+                    cmp.set_primary_key_maps(left_pk_map, right_pk_map)
+
                 res = cmp.compare(id, left, right)
                 if res:
                     findings.extend(res)