feat(relocation): Expand relocation model (#61106)

We add some useful features to the model, including:

- The ability to "pause" relocations before a certain step.
- The ability to cleanly "cancel" relocations before a certain step.
- A tracker for the last time we sent out the "unclaimed email" blast.

Issue: getsentry/team-ospo#222

migrations_lockfile.txt

@@ -9,5 +9,5 @@ feedback: 0003_feedback_add_env
 hybridcloud: 0009_make_user_id_optional_for_slug_reservation_replica
 nodestore: 0002_nodestore_no_dictfield
 replays: 0003_add_size_to_recording_segment
-sentry: 0611_add_regression_group_model
+sentry: 0612_expand_relocation_model
 social_auth: 0002_default_auto_field

src/sentry/migrations/0612_expand_relocation_model.py

@@ -0,0 +1,64 @@
+# Generated by Django 3.2.23 on 2023-12-05 01:30
+
+import django.db.models.expressions
+from django.db import migrations, models
+
+from sentry.new_migrations.migrations import CheckedMigration
+
+
+class Migration(CheckedMigration):
+    # This flag is used to mark that a migration shouldn't be automatically run in production. For
+    # the most part, this should only be used for operations where it's safe to run the migration
+    # after your code has deployed. So this should not be used for most operations that alter the
+    # schema of a table.
+    # Here are some things that make sense to mark as dangerous:
+    # - Large data migrations. Typically we want these to be run manually by ops so that they can
+    #   be monitored and not block the deploy for a long period of time while they run.
+    # - Adding indexes to large tables. Since this can take a long time, we'd generally prefer to
+    #   have ops run this and not block the deploy. Note that while adding an index is a schema
+    #   change, it's completely safe to run the operation after the code has deployed.
+    is_dangerous = False
+
+    dependencies = [
+        ("sentry", "0611_add_regression_group_model"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="relocation",
+            name="latest_unclaimed_emails_sent_at",
+            field=models.DateTimeField(default=None, null=True),
+        ),
+        migrations.AddField(
+            model_name="relocation",
+            name="scheduled_cancel_at_step",
+            field=models.SmallIntegerField(default=None, null=True),
+        ),
+        migrations.AddField(
+            model_name="relocation",
+            name="scheduled_pause_at_step",
+            field=models.SmallIntegerField(default=None, null=True),
+        ),
+        migrations.AddConstraint(
+            model_name="relocation",
+            constraint=models.CheckConstraint(
+                check=models.Q(
+                    ("scheduled_pause_at_step__gt", django.db.models.expressions.F("step")),
+                    ("scheduled_pause_at_step__isnull", True),
+                    _connector="OR",
+                ),
+                name="scheduled_pause_at_step_greater_than_current_step",
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="relocation",
+            constraint=models.CheckConstraint(
+                check=models.Q(
+                    ("scheduled_cancel_at_step__gt", django.db.models.expressions.F("step")),
+                    ("scheduled_cancel_at_step__isnull", True),
+                    _connector="OR",
+                ),
+                name="scheduled_cancel_at_step_greater_than_current_step",
+            ),
+        ),
+    ]

src/sentry/models/relocation.py

@@ -51,10 +51,30 @@ class Relocation(DefaultFieldsModel):
         def get_choices(cls) -> list[tuple[int, str]]:
             return [(key.value, key.name) for key in cls]
 
+        # Like `get_choices` above, except it excludes the final `COMPLETED` step.
+        @classmethod
+        def get_in_progress_choices(cls) -> list[tuple[int, str]]:
+            return [(key.value, key.name) for key in cls if key.name != "COMPLETED"]
+
+        @classmethod
+        def max_value(cls):
+            return max(item.value for item in cls)
+
     class Status(Enum):
         IN_PROGRESS = 0
         FAILURE = 1
         SUCCESS = 2
+        PAUSE = 3
+
+        # TODO(getsentry/team-ospo#190): Could we dedup this with a mixin in the future?
+        @classmethod
+        def get_choices(cls) -> list[tuple[int, str]]:
+            return [(key.value, key.name) for key in cls]
+
+    class EmailKind(Enum):
+        STARTED = 0
+        FAILED = 1
+        SUCCEEDED = 2
 
         # TODO(getsentry/team-ospo#190): Could we dedup this with a mixin in the future?
         @classmethod
@@ -84,6 +104,18 @@ class Relocation(DefaultFieldsModel):
         choices=Status.get_choices(), default=Status.IN_PROGRESS.value
     )
 
+    # Schedules a pause prior to some step that has not yet occurred. Useful to perform an orderly
+    # halting of the relocation. When unpausing, the unpausing process is responsible for scheduling
+    # the correct celery task so that the relocation may continue.
+    scheduled_pause_at_step = models.SmallIntegerField(
+        choices=Step.get_in_progress_choices(), null=True, default=None
+    )
+
+    # Schedules the termination of this relocation prior to some step that has not yet occurred.
+    scheduled_cancel_at_step = models.SmallIntegerField(
+        choices=Step.get_in_progress_choices(), null=True, default=None
+    )
+
     # Organizations, identified by slug, which this relocation seeks to import, specified by the
     # user as part of their relocation request. These slugs are relative to the import file, not
     # their final value post-import, which may be changed to avoid collisions.
@@ -96,7 +128,12 @@ class Relocation(DefaultFieldsModel):
 
     # The last status for which we have notified the user. It is `None` by default, to indicate that
     # we have not yet sent the user a "your relocation is in progress" email.
-    latest_notified = models.SmallIntegerField(choices=Step.get_choices(), null=True, default=None)
+    latest_notified = models.SmallIntegerField(
+        choices=EmailKind.get_choices(), null=True, default=None
+    )
+
+    # The last time we've sent the "claim your account" email blast to all unclaimed users.
+    latest_unclaimed_emails_sent_at = models.DateTimeField(null=True, default=None)
 
     # The last task started by this relocation. Because tasks for a given relocation are always
     # attempted sequentially, and never in concurrently (that is, there is always at most one task
@@ -114,6 +151,18 @@ class Relocation(DefaultFieldsModel):
     class Meta:
         app_label = "sentry"
         db_table = "sentry_relocation"
+        constraints = [
+            models.CheckConstraint(
+                name="scheduled_pause_at_step_greater_than_current_step",
+                check=models.Q(scheduled_pause_at_step__gt=models.F("step"))
+                | models.Q(scheduled_pause_at_step__isnull=True),
+            ),
+            models.CheckConstraint(
+                name="scheduled_cancel_at_step_greater_than_current_step",
+                check=models.Q(scheduled_cancel_at_step__gt=models.F("step"))
+                | models.Q(scheduled_cancel_at_step__isnull=True),
+            ),
+        ]
 
 
 @region_silo_only_model

src/sentry/tasks/relocation.py

@@ -51,7 +51,6 @@ from sentry.utils.env import gcp_project_id, log_gcp_credentials_details
 from sentry.utils.relocation import (
     RELOCATION_BLOB_SIZE,
     RELOCATION_FILE_TYPE,
-    EmailKind,
     LoggingPrinter,
     OrderedTask,
     create_cloudbuild_yaml,
@@ -323,7 +322,7 @@ def preprocessing_scan(uuid: str) -> None:
             # hours" email.
             send_relocation_update_email(
                 relocation,
-                EmailKind.STARTED,
+                Relocation.EmailKind.STARTED,
                 {
                     "uuid": str(relocation.uuid),
                     "orgs": relocation.want_org_slugs,
@@ -1140,7 +1139,7 @@ def notifying_owner(uuid: str) -> None:
     ):
         send_relocation_update_email(
             relocation,
-            EmailKind.SUCCEEDED,
+            Relocation.EmailKind.SUCCEEDED,
             {
                 "uuid": str(relocation.uuid),
                 "orgs": relocation.want_org_slugs,

src/sentry/utils/relocation.py

@@ -303,14 +303,8 @@ class LoggingPrinter(Printer):
             )
 
 
-class EmailKind(Enum):
-    STARTED = 0
-    FAILED = 1
-    SUCCEEDED = 2
-
-
 def send_relocation_update_email(
-    relocation: Relocation, email_kind: EmailKind, args: dict[str, Any]
+    relocation: Relocation, email_kind: Relocation.EmailKind, args: dict[str, Any]
 ):
     name = str(email_kind.name)
     name_lower = name.lower()
@@ -415,7 +409,7 @@ def fail_relocation(relocation: Relocation, task: OrderedTask, reason: str = "")
     logger.info("Task failed", extra={"uuid": relocation.uuid, "task": task.name, "reason": reason})
     send_relocation_update_email(
         relocation,
-        EmailKind.FAILED,
+        Relocation.EmailKind.FAILED,
         {
             "uuid": str(relocation.uuid),
             "reason": reason,