feat(backup): Add export checkpointer (#80711)

This feature mirrors what we do for importing, where we periodically
"save our work", so that if we experience an ephemeral failure (timeout,
pod restart, OOM, etc), we can "pick up where we left off". For imports,
we do this by saving `ImportChunk`s to the database every time we import
a few models, which allows us to check what we've already imported to
avoiding redoing work when retrying.

We use a similar strategy here for exporting. For every model kind, we
save a copy of the JSON of all instances of that model that we exported
to some durable media in specially-named "checkpoint" files. If there is
a failure midway through the export process, when we try again, we can
scan for these files to quickly re-use them, rather than doing very
expensive and resource intensive database queries again. While this does
assume that the model state has stayed relatively consistent between
runs, this is already an assumption we make for exporting in general (we
can't export a "single snapshot in time" of the database at once
anyway).

A follow-up PR will implement a subclass of `ExportCheckpointer` for
GCP, which is what we will use to checkpoint large SaaS->SaaS
relocations.

Author: azaslavsky

Diff for: src/sentry/backup/crypto.py

@@ -54,8 +54,6 @@ class Encryptor(ABC):
     A `IO[bytes]`-wrapper that contains relevant information and methods to encrypt some an in-memory JSON-ifiable dict.
     """
 
-    __fp: IO[bytes]
-
     @abstractmethod
     def get_public_key_pem(self) -> bytes:
         pass
@@ -67,10 +65,10 @@ class LocalFileEncryptor(Encryptor):
     """
 
     def __init__(self, fp: IO[bytes]):
-        self.__fp = fp
+        self.__key = fp.read()
 
     def get_public_key_pem(self) -> bytes:
-        return self.__fp.read()
+        return self.__key
 
 
 class GCPKMSEncryptor(Encryptor):
@@ -82,7 +80,7 @@ class GCPKMSEncryptor(Encryptor):
     crypto_key_version: CryptoKeyVersion | None = None
 
     def __init__(self, fp: IO[bytes]):
-        self.__fp = fp
+        self.__key = fp.read()
 
     @classmethod
     def from_crypto_key_version(cls, crypto_key_version: CryptoKeyVersion) -> GCPKMSEncryptor:
@@ -93,7 +91,7 @@ def from_crypto_key_version(cls, crypto_key_version: CryptoKeyVersion) -> GCPKMS
     def get_public_key_pem(self) -> bytes:
         if self.crypto_key_version is None:
             # Read the user supplied configuration into the proper format.
-            gcp_kms_config_json = orjson.loads(self.__fp.read())
+            gcp_kms_config_json = orjson.loads(self.__key)
             try:
                 self.crypto_key_version = CryptoKeyVersion(**gcp_kms_config_json)
             except TypeError:
@@ -217,8 +215,6 @@ class Decryptor(ABC):
     tarball.
     """
 
-    __fp: IO[bytes]
-
     @abstractmethod
     def read(self) -> bytes:
         pass
@@ -234,22 +230,22 @@ class LocalFileDecryptor(Decryptor):
     """
 
     def __init__(self, fp: IO[bytes]):
-        self.__fp = fp
+        self.__key = fp.read()
 
     @classmethod
     def from_bytes(cls, b: bytes) -> LocalFileDecryptor:
         return cls(io.BytesIO(b))
 
     def read(self) -> bytes:
-        return self.__fp.read()
+        return self.__key
 
     def decrypt_data_encryption_key(self, unwrapped: UnwrappedEncryptedExportTarball) -> bytes:
         """
         Decrypt the encrypted data encryption key used to encrypt the actual export JSON.
         """
 
         # Compare the public and private key, to ensure that they are a match.
-        private_key_pem = self.__fp.read()
+        private_key_pem = self.__key
         private_key = serialization.load_pem_private_key(
             private_key_pem,
             password=None,
@@ -286,17 +282,17 @@ class GCPKMSDecryptor(Decryptor):
     """
 
     def __init__(self, fp: IO[bytes]):
-        self.__fp = fp
+        self.__key = fp.read()
 
     @classmethod
     def from_bytes(cls, b: bytes) -> GCPKMSDecryptor:
         return cls(io.BytesIO(b))
 
     def read(self) -> bytes:
-        return self.__fp.read()
+        return self.__key
 
     def decrypt_data_encryption_key(self, unwrapped: UnwrappedEncryptedExportTarball) -> bytes:
-        gcp_kms_config_bytes = self.__fp.read()
+        gcp_kms_config_bytes = self.__key
 
         # Read the user supplied configuration into the proper format.
         gcp_kms_config_json = orjson.loads(gcp_kms_config_bytes)
@@ -345,3 +341,13 @@ def decrypt_encrypted_tarball(tarball: IO[bytes], decryptor: Decryptor) -> bytes
     decrypted_dek = decryptor.decrypt_data_encryption_key(unwrapped)
     fernet = Fernet(decrypted_dek)
     return fernet.decrypt(unwrapped.encrypted_json_blob)
+
+
+class EncryptorDecryptorPair:
+    """
+    An Encryptor and Decryptor that use paired public and private keys, respectively.
+    """
+
+    def __init__(self, encryptor: Encryptor, decryptor: Decryptor):
+        self.encryptor = encryptor
+        self.decryptor = decryptor

Diff for: src/sentry/backup/exports.py

@@ -5,12 +5,15 @@
 # We have to use the default JSON interface to enable pretty-printing on export. When loading JSON,
 # we still use the one from `sentry.utils`, imported as `sentry_json` below.
 import json as builtin_json  # noqa: S003
+from abc import ABC, abstractmethod
 from typing import IO
 
 import orjson
 
-from sentry.backup.crypto import Encryptor, create_encrypted_export_tarball
+from sentry.backup.crypto import Encryptor, EncryptorDecryptorPair, create_encrypted_export_tarball
 from sentry.backup.dependencies import (
+    ImportKind,
+    NormalizedModelName,
     PrimaryKeyMap,
     dependencies,
     get_model_name,
@@ -20,6 +23,7 @@
 from sentry.backup.scopes import ExportScope
 from sentry.backup.services.import_export.model import (
     RpcExportError,
+    RpcExportOk,
     RpcExportScope,
     RpcFilter,
     RpcPrimaryKeyMap,
@@ -41,6 +45,69 @@ def __init__(self, context: RpcExportError) -> None:
         self.context = context
 
 
+class ExportCheckpointerError(Exception):
+    pass
+
+
+class ExportCheckpointer(ABC):
+    """
+    For very large exports, the exporting environment may fall over half-way through the process:
+    the thread running it may hit some timeout, or it may OOM, or fail for some other ephemeral
+    reason. To help in such situations, we'd like an API for saving "checkpoints" during the export.
+
+    This class provides per-model checkpointing support for exports. Since there is a topologically
+    sorted order of models being exported, as we move through this list, we can save the exported
+    JSON for each kind of model in order to some stable media (disk, GCP, etc). If there is a
+    failure late in the export process, when it is retried, the exporter can check if that
+    particular model already exists in the checkpointer's cache, thereby avoiding redoing the work
+    of pulling the models from the database, processing them, etc. This ensures that in most retry
+    situations, we can quickly "re-ingest" already-exported models in memory and pick up where we
+    left off.
+    """
+
+    def _parse_cached_json(self, json_data: bytes) -> RpcExportOk | None:
+        max_pk = 0
+        pk_map = PrimaryKeyMap()
+        models = orjson.loads(json_data)
+        for model in models:
+            model_name = model.get("model", None)
+            pk = model.get("pk", None)
+            if model_name is None or pk is None:
+                raise ExportCheckpointerError("Improperly formatted entry")
+
+            pk_map.insert(model_name, pk, pk, ImportKind.Inserted)
+            if pk > max_pk:
+                max_pk = pk
+
+        return RpcExportOk(
+            mapped_pks=RpcPrimaryKeyMap.into_rpc(pk_map), max_pk=max_pk, json_data=json_data
+        )
+
+    @abstractmethod
+    def get(self, model_name: NormalizedModelName) -> RpcExportOk | None:
+        pass
+
+    @abstractmethod
+    def add(self, model_name: NormalizedModelName, json_data: str) -> None:
+        pass
+
+
+class NoopExportCheckpointer(ExportCheckpointer):
+    """
+    A noop checkpointer - that is, it doesn't write or read any checkpoints, always returning None.
+    This means that no checkpointing ever occurs.
+    """
+
+    def __init__(self, crypto: EncryptorDecryptorPair | None, printer: Printer):
+        pass
+
+    def get(self, model_name: NormalizedModelName) -> RpcExportOk | None:
+        return None
+
+    def add(self, model_name: NormalizedModelName, json_data: str) -> None:
+        return None
+
+
 def _export(
     dest: IO[bytes],
     scope: ExportScope,
@@ -49,6 +116,7 @@ def _export(
     indent: int = 2,
     filter_by: Filter | None = None,
     printer: Printer,
+    checkpointer: ExportCheckpointer | None = None,
 ):
     """
     Exports core data for the Sentry installation.
@@ -68,6 +136,7 @@ def _export(
         printer.echo(errText, err=True)
         raise RuntimeError(errText)
 
+    cache = checkpointer if checkpointer is not None else NoopExportCheckpointer(None, printer)
     json_export = []
     pk_map = PrimaryKeyMap()
     allowed_relocation_scopes = scope.value
@@ -119,25 +188,33 @@ def _export(
 
         dep_models = {get_model_name(d) for d in model_relations.get_dependencies_for_relocation()}
         export_by_model = ImportExportService.get_exporter_for_model(model)
-        result = export_by_model(
-            export_model_name=str(model_name),
-            scope=RpcExportScope.into_rpc(scope),
-            from_pk=0,
-            filter_by=[RpcFilter.into_rpc(f) for f in filters],
-            pk_map=RpcPrimaryKeyMap.into_rpc(pk_map.partition(dep_models)),
-            indent=indent,
+        cached_result = cache.get(model_name)
+        result = (
+            cached_result
+            if cached_result is not None
+            else export_by_model(
+                export_model_name=str(model_name),
+                scope=RpcExportScope.into_rpc(scope),
+                from_pk=0,
+                filter_by=[RpcFilter.into_rpc(f) for f in filters],
+                pk_map=RpcPrimaryKeyMap.into_rpc(pk_map.partition(dep_models)),
+                indent=indent,
+            )
         )
 
         if isinstance(result, RpcExportError):
             printer.echo(result.pretty(), err=True)
             raise ExportingError(result)
 
         pk_map.extend(result.mapped_pks.from_rpc())
+        json_models = orjson.loads(result.json_data)
+        if cached_result is None:
+            cache.add(model_name, json_models)
 
         # TODO(getsentry/team-ospo#190): Since the structure of this data is very predictable (an
         # array of serialized model objects), we could probably avoid re-ingesting the JSON string
         # as a future optimization.
-        for json_model in orjson.loads(result.json_data):
+        for json_model in json_models:
             json_export.append(json_model)
 
     # If no `encryptor` argument was passed in, this is an unencrypted export, so we can just dump
@@ -158,6 +235,7 @@ def export_in_user_scope(
     user_filter: set[str] | None = None,
     indent: int = 2,
     printer: Printer,
+    checkpointer: ExportCheckpointer | None = None,
 ):
     """
     Perform an export in the `User` scope, meaning that only models with `RelocationScope.User` will
@@ -174,6 +252,7 @@ def export_in_user_scope(
         filter_by=Filter(User, "username", user_filter) if user_filter is not None else None,
         indent=indent,
         printer=printer,
+        checkpointer=checkpointer,
     )
 
 
@@ -184,6 +263,7 @@ def export_in_organization_scope(
     org_filter: set[str] | None = None,
     indent: int = 2,
     printer: Printer,
+    checkpointer: ExportCheckpointer | None = None,
 ):
     """
     Perform an export in the `Organization` scope, meaning that only models with
@@ -201,6 +281,7 @@ def export_in_organization_scope(
         filter_by=Filter(Organization, "slug", org_filter) if org_filter is not None else None,
         indent=indent,
         printer=printer,
+        checkpointer=checkpointer,
     )
 
 
@@ -210,6 +291,7 @@ def export_in_config_scope(
     encryptor: Encryptor | None = None,
     indent: int = 2,
     printer: Printer,
+    checkpointer: ExportCheckpointer | None = None,
 ):
     """
     Perform an export in the `Config` scope, meaning that only models directly related to the global
@@ -226,6 +308,7 @@ def export_in_config_scope(
         filter_by=Filter(User, "pk", import_export_service.get_all_globally_privileged_users()),
         indent=indent,
         printer=printer,
+        checkpointer=checkpointer,
     )
 
 
@@ -235,6 +318,7 @@ def export_in_global_scope(
     encryptor: Encryptor | None = None,
     indent: int = 2,
     printer: Printer,
+    checkpointer: ExportCheckpointer | None = None,
 ):
     """
     Perform an export in the `Global` scope, meaning that all models will be exported from the
@@ -246,4 +330,5 @@ def export_in_global_scope(
         encryptor=encryptor,
         indent=indent,
         printer=printer,
+        checkpointer=checkpointer,
     )