Import: Support import of huddles.

zerver/lib/import_realm.py

@@ -16,6 +16,7 @@
 from zerver.lib.avatar_hash import user_avatar_path_from_ids
 from zerver.lib.bulk_create import bulk_create_users
 from zerver.lib.create_user import random_api_key
+from zerver.lib.utils import make_safe_digest
 from zerver.lib.export import DATE_FIELDS, realm_tables, \
     Record, TableData, TableName, Field, Path
 from zerver.lib.upload import random_name, sanitize_name, \
@@ -25,7 +26,7 @@
     RealmDomain, Recipient, get_user_profile_by_id, \
     UserPresence, UserActivity, UserActivityInterval, Reaction, \
     CustomProfileField, CustomProfileFieldValue, \
-    Attachment, get_system_bot, email_to_username
+    Attachment, get_system_bot, email_to_username, get_huddle_hash
 
 # Code from here is the realm import code path
 
@@ -41,6 +42,7 @@
 id_maps = {
     'client': {},
     'user_profile': {},
+    'huddle': {},
     'realm': {},
     'stream': {},
     'recipient': {},
@@ -58,8 +60,13 @@
     'customprofilefield': {},
     'customprofilefield_value': {},
     'attachment': {},
+    'recipient_to_huddle_map': {},
 }  # type: Dict[str, Dict[int, int]]
 
+id_map_to_list = {
+    'huddle_to_user_list': {},
+}  # type: Dict[str, Dict[int, List[int]]]
+
 path_maps = {
     'attachment_path': {},
 }  # type: Dict[str, Dict[str, str]]
@@ -94,6 +101,27 @@ def fix_upload_links(data: TableData, message_table: TableName) -> None:
                     if message['rendered_content']:
                         message['rendered_content'] = message['rendered_content'].replace(key, value)
 
+def process_huddle_hash(data: TableData, table: TableName) -> None:
+    """
+    Build new huddle hashes with the updated ids of the users
+    """
+    for huddle in data[table]:
+        user_id_list = id_map_to_list['huddle_to_user_list'][huddle['id']]
+        huddle['huddle_hash'] = get_huddle_hash(user_id_list)
+
+def get_huddles_from_subscription(data: TableData, table: TableName) -> None:
+    """
+    Extract the user_profiles involved in a huddle from the subscription object
+    This helps to generate a unique huddle hash from the updated user_profile ids
+    """
+    for key, value in id_maps['recipient_to_huddle_map'].items():
+        id_map_to_list['huddle_to_user_list'][value] = []
+
+    for subscription in data[table]:
+        if subscription['recipient'] in id_maps['recipient_to_huddle_map'].keys():
+            huddle_id = id_maps['recipient_to_huddle_map'][subscription['recipient']]
+            id_map_to_list['huddle_to_user_list'][huddle_id].append(subscription['user_profile_id'])
+
 def current_table_ids(data: TableData, table: TableName) -> List[int]:
     """
     Returns the ids present in the current table
@@ -169,14 +197,21 @@ def re_map_foreign_keys_internal(data_table: List[Record],
     '''
     lookup_table = id_maps[related_table]
     for item in data_table:
+        old_id = item[field_name]
         if recipient_field:
             if related_table == "stream" and item['type'] == 2:
                 pass
             elif related_table == "user_profile" and item['type'] == 1:
                 pass
+            elif related_table == "huddle" and item['type'] == 3:
+                # save the recipient id with the huddle id, so that we can extract
+                # the user_profile ids involved in a huddle with the help of the
+                # subscription object
+                # check function 'get_huddles_from_subscription'
+                id_maps['recipient_to_huddle_map'][item['id']] = lookup_table[old_id]
+                pass
             else:
                 continue
-        old_id = item[field_name]
         if old_id in lookup_table:
             new_id = lookup_table[old_id]
             if verbose:
@@ -481,20 +516,27 @@ def do_import_realm(import_dir: Path, subdomain: str) -> Realm:
     UserProfile.objects.bulk_create(user_profiles)
 
     if 'zerver_huddle' in data:
-        bulk_import_model(data, Huddle, 'zerver_huddle')
+        update_model_ids(Huddle, data, 'zerver_huddle', 'huddle')
 
     re_map_foreign_keys(data, 'zerver_recipient', 'type_id', related_table="stream",
                         recipient_field=True, id_field=True)
     re_map_foreign_keys(data, 'zerver_recipient', 'type_id', related_table="user_profile",
                         recipient_field=True, id_field=True)
+    re_map_foreign_keys(data, 'zerver_recipient', 'type_id', related_table="huddle",
+                        recipient_field=True, id_field=True)
     update_model_ids(Recipient, data, 'zerver_recipient', 'recipient')
     bulk_import_model(data, Recipient, 'zerver_recipient')
 
     re_map_foreign_keys(data, 'zerver_subscription', 'user_profile', related_table="user_profile")
+    get_huddles_from_subscription(data, 'zerver_subscription')
     re_map_foreign_keys(data, 'zerver_subscription', 'recipient', related_table="recipient")
     update_model_ids(Subscription, data, 'zerver_subscription', 'subscription')
     bulk_import_model(data, Subscription, 'zerver_subscription')
 
+    if 'zerver_huddle' in data:
+        process_huddle_hash(data, 'zerver_huddle')
+        bulk_import_model(data, Huddle, 'zerver_huddle')
+
     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')