import script: re_map minor foreign keys.

zerver/lib/export.py

@@ -53,10 +53,10 @@ PostProcessData = Any  # TODO: make more specific
 # TODO: tighten this up with a union.
 MessageOutput = Dict[str, Any]
 
-realm_tables = [("zerver_defaultstream", DefaultStream),
-                ("zerver_realmemoji", RealmEmoji),
-                ("zerver_realmdomain", RealmDomain),
-                ("zerver_realmfilter", RealmFilter)]  # List[Tuple[TableName, Any]]
+realm_tables = [("zerver_defaultstream", DefaultStream, "defaultstream"),
+                ("zerver_realmemoji", RealmEmoji, "realmemoji"),
+                ("zerver_realmdomain", RealmDomain, "realmdomain"),
+                ("zerver_realmfilter", RealmFilter, "realmfilter")]  # List[Tuple[TableName, Any, str]]
 
 
 ALL_ZERVER_TABLES = [
@@ -1222,7 +1222,15 @@ id_maps = {
     'stream': {},
     'recipient': {},
     'subscription': {},
+    'defaultstream': {},
+    'realmemoji': {},
+    'realmdomain': {},
+    'realmfilter': {},
     'message': {},
+    'user_presence': {},
+    'useractivity': {},
+    'useractivityinterval': {},
+    'usermessage': {},
 }  # type: Dict[str, Dict[int, int]]
 
 def update_id_map(table: TableName, old_id: int, new_id: int) -> None:
@@ -1250,6 +1258,8 @@ def current_table_ids(data: TableData, table: TableName) -> List[int]:
     return id_list
 
 def idseq(model_class: Any) -> str:
+    if model_class == RealmDomain:
+        return 'zerver_realmalias_id_seq'
     return '{}_id_seq'.format(model_class._meta.db_table)
 
 def allocate_ids(model_class: Any, count: int) -> List[int]:
@@ -1515,8 +1525,14 @@ def do_import_realm(import_dir: Path) -> Realm:
     realm.save()
 
     re_map_foreign_keys(data['zerver_defaultstream'], 'stream', related_table="stream")
-    for (table, model) in realm_tables:
+    re_map_foreign_keys(data['zerver_realmemoji'], 'author', related_table="user_profile")
+    for (table, model, related_table) in realm_tables:
         re_map_foreign_keys(data[table], 'realm', related_table="realm")
+        old_id_list = current_table_ids(data, table)
+        allocated_id_list = allocate_ids(model, len(data[table]))
+        for item in range(len(data[table])):
+            update_id_map(related_table, old_id_list[item], allocated_id_list[item])
+        re_map_foreign_keys(data[table], 'id', related_table=related_table, id_field=True)
         bulk_import_model(data, model, table)
 
     # Remap the user IDs for notification_bot and friends to their
@@ -1582,15 +1598,35 @@ def do_import_realm(import_dir: Path) -> Realm:
     fix_datetime_fields(data, 'zerver_userpresence')
     re_map_foreign_keys(data['zerver_userpresence'], 'user_profile', related_table="user_profile")
     re_map_foreign_keys(data['zerver_userpresence'], 'client', related_table='client')
+    userpresence_id_list = current_table_ids(data, 'zerver_userpresence')
+    allocated_userpresence_id_list = allocate_ids(UserPresence, len(data['zerver_userpresence']))
+    for item in range(len(data['zerver_userpresence'])):
+        update_id_map('user_presence', userpresence_id_list[item],
+                      allocated_userpresence_id_list[item])
+    re_map_foreign_keys(data['zerver_userpresence'], 'id', related_table="user_presence", id_field=True)
     bulk_import_model(data, UserPresence, 'zerver_userpresence')
 
     fix_datetime_fields(data, 'zerver_useractivity')
     re_map_foreign_keys(data['zerver_useractivity'], 'user_profile', related_table="user_profile")
     re_map_foreign_keys(data['zerver_useractivity'], 'client', related_table='client')
+    useractivity_id_list = current_table_ids(data, 'zerver_useractivity')
+    allocated_useractivity_id_list = allocate_ids(UserActivity, len(data['zerver_useractivity']))
+    for item in range(len(data['zerver_useractivity'])):
+        update_id_map('useractivity', useractivity_id_list[item],
+                      allocated_useractivity_id_list[item])
+    re_map_foreign_keys(data['zerver_useractivity'], 'id', related_table="useractivity", id_field=True)
     bulk_import_model(data, UserActivity, 'zerver_useractivity')
 
     fix_datetime_fields(data, 'zerver_useractivityinterval')
     re_map_foreign_keys(data['zerver_useractivityinterval'], 'user_profile', related_table="user_profile")
+    useractivityinterval_id_list = current_table_ids(data, 'zerver_useractivityinterval')
+    allocated_useractivityinterval_id_list = allocate_ids(UserActivityInterval,
+                                                          len(data['zerver_useractivityinterval']))
+    for item in range(len(data['zerver_useractivityinterval'])):
+        update_id_map('useractivityinterval', useractivityinterval_id_list[item],
+                      allocated_useractivityinterval_id_list[item])
+    re_map_foreign_keys(data['zerver_useractivityinterval'], 'id', related_table="useractivityinterval",
+                        id_field=True)
     bulk_import_model(data, UserActivityInterval, 'zerver_useractivityinterval')
 
     # Import uploaded files and avatars
@@ -1661,6 +1697,13 @@ def import_message_data(import_dir: Path) -> None:
         re_map_foreign_keys(data['zerver_usermessage'], 'message', related_table="message")
         re_map_foreign_keys(data['zerver_usermessage'], 'user_profile', related_table="user_profile")
         fix_bitfield_keys(data, 'zerver_usermessage', 'flags')
+        usermessage_id_list = current_table_ids(data, 'zerver_usermessage')
+        allocated_usermessage_id_list = allocate_ids(UserMessage, len(data['zerver_usermessage']))
+        for item in range(len(data['zerver_usermessage'])):
+            update_id_map('usermessage', usermessage_id_list[item],
+                          allocated_usermessage_id_list[item])
+        re_map_foreign_keys(data['zerver_usermessage'], 'id', related_table="usermessage",
+                            id_field=True)
         bulk_import_model(data, UserMessage, 'zerver_usermessage')
 
         dump_file_id += 1