import: Extract re_map_realm_emoji_codes.

zerver/lib/import_realm.py

@@ -427,7 +427,6 @@ def re_map_foreign_keys(
     verbose: bool = False,
     id_field: bool = False,
     recipient_field: bool = False,
-    reaction_field: bool = False,
 ) -> None:
     """
     This is a wrapper function for all the realm data tables
@@ -447,7 +446,6 @@ def re_map_foreign_keys(
         verbose,
         id_field,
         recipient_field,
-        reaction_field,
     )
 
 
@@ -459,7 +457,6 @@ def re_map_foreign_keys_internal(
     verbose: bool = False,
     id_field: bool = False,
     recipient_field: bool = False,
-    reaction_field: bool = False,
 ) -> None:
     """
     We occasionally need to assign new ids to rows during the
@@ -487,11 +484,6 @@ def re_map_foreign_keys_internal(
             else:
                 continue
         old_id = item[field_name]
-        if reaction_field:
-            if item["reaction_type"] == Reaction.REALM_EMOJI:
-                old_id = int(old_id)
-            else:
-                continue
         if old_id in lookup_table:
             new_id = lookup_table[old_id]
             if verbose:
@@ -503,13 +495,39 @@ def re_map_foreign_keys_internal(
         if not id_field:
             item[field_name + "_id"] = new_id
             del item[field_name]
-        else:
-            if reaction_field:
-                item[field_name] = str(new_id)
         else:
             item[field_name] = new_id
 
 
+def re_map_realm_emoji_codes(data: TableData, *, table_name: str) -> None:
+    """
+    Some tables, including Reaction and UserStatus, contain a form of
+    foreign key reference to the RealmEmoji table in the form of
+    `str(realm_emoji.id)` when `reaction_type="realm_emoji"`.
+
+    See the block comment for emoji_code in the AbstractEmoji
+    definition for more details.
+    """
+    realm_emoji_dct = {}
+
+    for row in data["zerver_realmemoji"]:
+        realm_emoji_dct[row["id"]] = row
+
+    for row in data[table_name]:
+        if row["reaction_type"] == Reaction.REALM_EMOJI:
+            old_realm_emoji_id = int(row["emoji_code"])
+
+            # Fail hard here if we didn't map correctly here
+            new_realm_emoji_id = ID_MAP["realmemoji"][old_realm_emoji_id]
+
+            # This is a very important sanity check.
+            realm_emoji_row = realm_emoji_dct[new_realm_emoji_id]
+            assert realm_emoji_row["name"] == row["emoji_name"]
+
+            # Now update emoji_code to the new id.
+            row["emoji_code"] = str(new_realm_emoji_id)
+
+
 def re_map_foreign_keys_many_to_many(
     data: TableData,
     table: TableName,
@@ -1228,14 +1246,7 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea
 
     re_map_foreign_keys(data, "zerver_reaction", "message", related_table="message")
     re_map_foreign_keys(data, "zerver_reaction", "user_profile", related_table="user_profile")
-    re_map_foreign_keys(
-        data,
-        "zerver_reaction",
-        "emoji_code",
-        related_table="realmemoji",
-        id_field=True,
-        reaction_field=True,
-    )
+    re_map_realm_emoji_codes(data, table_name="zerver_reaction")
     update_model_ids(Reaction, data, "reaction")
     bulk_import_model(data, Reaction)