feat(relocation): Add import org and user ids to GET /relocations/:id (#69280)

This will helps us render this information in the admin panel, which
will in turn make it much easier to see what actually happened during a
relocation.

src/sentry/api/serializers/models/relocation.py

@@ -1,7 +1,10 @@
+import dataclasses
 from collections.abc import Mapping, MutableMapping, Sequence
 from typing import Any
 
 from sentry.api.serializers import Serializer, register
+from sentry.db.models.manager.base import BaseManager
+from sentry.models.importchunk import BaseImportChunk, ControlImportChunkReplica, RegionImportChunk
 from sentry.models.relocation import Relocation
 from sentry.models.user import User
 from sentry.services.hybrid_cloud.user.model import RpcUser
@@ -9,10 +12,42 @@ from sentry.services.hybrid_cloud.user.service import user_service
 from sentry.utils.json import JSONData
 
 
+@dataclasses.dataclass(frozen=True)
+class RelocationMetadata:
+    """
+    Some useful info to collect about a relocation when serving it.
+    """
+
+    # Maps the creator/owner's (aka "meta" users) `id`s to their respective `username`.
+    meta_users: Mapping[int, RpcUser]
+
+    # List the ids of the imported `User` models.
+    imported_user_ids: list[int]
+
+    # List the ids of the imported `Organization` models.
+    imported_org_ids: list[int]
+
+
+def get_all_imported_ids_of_model(chunks: BaseManager[BaseImportChunk]) -> list[int]:
+    all_imported_ids = set()
+    for chunk in chunks:
+        all_imported_ids |= (
+            set(chunk.inserted_map.values())
+            | set(chunk.existing_map.values())
+            | set(chunk.overwrite_map.values())
+        )
+
+    return list(all_imported_ids)
+
+
 @register(Relocation)
 class RelocationSerializer(Serializer):
     def serialize(
-        self, obj: Relocation, attrs: Mapping[int, RpcUser], user: User, **kwargs: Any
+        self,
+        obj: Relocation,
+        attrs: Any,
+        user: User,
+        **kwargs: Any,
     ) -> Mapping[str, JSONData]:
         scheduled_at_pause_step = (
             Relocation.Step(obj.scheduled_pause_at_step).name
@@ -34,12 +69,12 @@ class RelocationSerializer(Serializer):
             "dateAdded": obj.date_added,
             "dateUpdated": obj.date_updated,
             "uuid": str(obj.uuid),
-            "creatorEmail": attrs[obj.creator_id].email,
+            "creatorEmail": attrs.meta_users[obj.creator_id].email,
             "creatorId": str(obj.creator_id),
-            "creatorUsername": attrs[obj.creator_id].username,
-            "ownerEmail": attrs[obj.owner_id].email,
+            "creatorUsername": attrs.meta_users[obj.creator_id].username,
+            "ownerEmail": attrs.meta_users[obj.owner_id].email,
             "ownerId": str(obj.owner_id),
-            "ownerUsername": attrs[obj.owner_id].username,
+            "ownerUsername": attrs.meta_users[obj.owner_id].username,
             "status": Relocation.Status(obj.status).name,
             "step": Relocation.Step(obj.step).name,
             "failureReason": obj.failure_reason,
@@ -49,9 +84,11 @@ class RelocationSerializer(Serializer):
             "wantUsernames": obj.want_usernames,
             "latestNotified": latest_notified,
             "latestUnclaimedEmailsSentAt": obj.latest_unclaimed_emails_sent_at,
+            "importedUserIds": attrs.imported_user_ids,
+            "importedOrgIds": attrs.imported_org_ids,
         }
 
-    def get_attrs(
+    def get_attrs_old(
         self, item_list: Sequence[Relocation], user: User, **kwargs: Any
     ) -> MutableMapping[Relocation, Mapping[int, RpcUser]]:
         user_ids = set()
@@ -62,3 +99,32 @@ class RelocationSerializer(Serializer):
         users = user_service.get_many(filter=dict(user_ids=list(user_ids)))
         user_map = {u.id: u for u in users}
         return {r: user_map for r in item_list}
+
+    def get_attrs(
+        self, item_list: Sequence[Relocation], user: User, **kwargs: Any
+    ) -> MutableMapping[Relocation, RelocationMetadata]:
+        metadata_map = {}
+        for relocation in item_list:
+            # Pull down the "meta" users to finish building out the `RelocationMetadata`.
+            user_ids = set()
+            user_ids.add(relocation.creator_id)
+            user_ids.add(relocation.owner_id)
+            users = user_service.get_many(filter=dict(user_ids=list(user_ids)))
+
+            # If the Relocation has finished, we should be able to pull the imported user and org
+            # pks from the correct `ControlImportChunkReplica` and `RegionImportChunk` entries,
+            # respectively. If it has not yet completed the `IMPORTING` step, these will both be
+            # empty.
+            user_import_chunks = ControlImportChunkReplica.objects.filter(
+                import_uuid=str(relocation.uuid), model="sentry.user"
+            )
+            org_import_chunks = RegionImportChunk.objects.filter(
+                import_uuid=str(relocation.uuid), model="sentry.organization"
+            )
+            metadata_map[relocation] = RelocationMetadata(
+                meta_users={u.id: u for u in users},
+                imported_user_ids=get_all_imported_ids_of_model(user_import_chunks),
+                imported_org_ids=get_all_imported_ids_of_model(org_import_chunks),
+            )
+
+        return metadata_map

tests/sentry/api/endpoints/relocations/test_retry.py

@@ -110,6 +110,8 @@ class RetryRelocationTest(APITestCase):
         assert response.data["latestUnclaimedEmailsSentAt"] is None
         assert response.data["scheduledPauseAtStep"] is None
         assert response.data["wantUsernames"] is None
+        assert response.data["importedUserIds"] == []
+        assert response.data["importedOrgIds"] == []
 
         assert (
             Relocation.objects.filter(owner_id=self.owner.id)

tests/sentry/api/serializers/test_relocation.py

@@ -1,6 +1,7 @@
 from datetime import datetime, timezone
 
 from sentry.api.serializers import serialize
+from sentry.models.importchunk import ControlImportChunkReplica, RegionImportChunk
 from sentry.models.relocation import Relocation
 from sentry.testutils.cases import TestCase
 from sentry.testutils.helpers.datetime import freeze_time
@@ -22,6 +23,37 @@ class RelocationSerializerTest(TestCase):
         )
         self.login_as(user=self.superuser, superuser=True)
 
+        self.first_imported_user = self.create_user(email="first@example.com")
+        self.second_imported_user = self.create_user(email="second@example.com")
+        self.imported_org = self.create_organization(owner=self.first_imported_user)
+        self.create_member(
+            user=self.second_imported_user, organization=self.imported_org, role="member", teams=[]
+        )
+
+    def mock_imported_users_and_org(self, relocation: Relocation) -> None:
+        ControlImportChunkReplica.objects.create(
+            import_uuid=relocation.uuid,
+            model="sentry.user",
+            min_ordinal=1,
+            max_ordinal=2,
+            min_source_pk=1,
+            max_source_pk=2,
+            min_inserted_pk=1,
+            max_inserted_pk=2,
+            inserted_map={1: self.first_imported_user.id, 2: self.second_imported_user.id},
+        )
+        RegionImportChunk.objects.create(
+            import_uuid=relocation.uuid,
+            model="sentry.organization",
+            min_ordinal=1,
+            max_ordinal=2,
+            min_source_pk=1,
+            max_source_pk=2,
+            min_inserted_pk=1,
+            max_inserted_pk=2,
+            inserted_map={1: self.imported_org.id},
+        )
+
     def test_in_progress(self):
         relocation: Relocation = Relocation.objects.create(
             date_added=TEST_DATE_ADDED,
@@ -58,6 +90,8 @@ class RelocationSerializerTest(TestCase):
         assert not result["latestUnclaimedEmailsSentAt"]
         assert "latestTask" not in result
         assert "latestTaskAttempts" not in result
+        assert result["importedUserIds"] == []
+        assert result["importedOrgIds"] == []
 
     def test_pause(self):
         relocation: Relocation = Relocation.objects.create(
@@ -72,6 +106,7 @@ class RelocationSerializerTest(TestCase):
             latest_task=OrderedTask.IMPORTING.name,
             latest_task_attempts=1,
         )
+        self.mock_imported_users_and_org(relocation)
         result = serialize(relocation)
 
         assert result["dateAdded"] == TEST_DATE_ADDED
@@ -94,6 +129,11 @@ class RelocationSerializerTest(TestCase):
         assert not result["latestUnclaimedEmailsSentAt"]
         assert "latestTask" not in result
         assert "latestTaskAttempts" not in result
+        assert sorted(result["importedUserIds"]) == [
+            self.first_imported_user.id,
+            self.second_imported_user.id,
+        ]
+        assert result["importedOrgIds"] == [self.imported_org.id]
 
     def test_success(self):
         relocation: Relocation = Relocation.objects.create(
@@ -109,6 +149,7 @@ class RelocationSerializerTest(TestCase):
             latest_task=OrderedTask.COMPLETED.name,
             latest_task_attempts=1,
         )
+        self.mock_imported_users_and_org(relocation)
         result = serialize(relocation)
 
         assert result["dateAdded"] == TEST_DATE_ADDED
@@ -131,6 +172,11 @@ class RelocationSerializerTest(TestCase):
         assert result["latestUnclaimedEmailsSentAt"] == TEST_DATE_UPDATED
         assert "latestTask" not in result
         assert "latestTaskAttempts" not in result
+        assert sorted(result["importedUserIds"]) == [
+            self.first_imported_user.id,
+            self.second_imported_user.id,
+        ]
+        assert result["importedOrgIds"] == [self.imported_org.id]
 
     def test_failure(self):
         relocation: Relocation = Relocation.objects.create(
@@ -169,3 +215,5 @@ class RelocationSerializerTest(TestCase):
         assert not result["latestUnclaimedEmailsSentAt"]
         assert "latestTask" not in result
         assert "latestTaskAttempts" not in result
+        assert result["importedUserIds"] == []
+        assert result["importedOrgIds"] == []