From a2ecdeb28d0d671bf777c14f23076ee9fa3c275e Mon Sep 17 00:00:00 2001 From: Rhea Parekh Date: Tue, 13 Mar 2018 17:19:17 +0530 Subject: [PATCH] import script: re_map minor foreign keys. --- zerver/lib/export.py | 53 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 48 insertions(+), 5 deletions(-) diff --git a/zerver/lib/export.py b/zerver/lib/export.py index 20f44863e1..efd9d2d30c 100644 --- a/zerver/lib/export.py +++ b/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