import: Fix import of Stream.creator field.

When importing a Stream, UserProfiles don't yet exist. so trying to
import stream.creator fails with something like:

```
psycopg2.errors.ForeignKeyViolation: insert or update on table
"zerver_stream" violates foreign key constraint
"zerver_stream_creator_id_65aeba7e_fk_zerver_userprofile_id"
DETAIL:  Key (creator_id)=(5) is not present in table "zerver_userprofile".
```
This commit is contained in:
Mateusz Mandera 2024-05-14 03:40:42 +02:00 committed by Tim Abbott
parent 2e8fdab896
commit 160076fdfa
2 changed files with 28 additions and 0 deletions

View File

@ -1081,7 +1081,18 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea
# Stream objects are created by Django.
fix_datetime_fields(data, "zerver_stream")
re_map_foreign_keys(data, "zerver_stream", "realm", related_table="realm")
re_map_foreign_keys(data, "zerver_stream", "creator", related_table="user_profile")
# There's a circular dependency between Stream and UserProfile due to
# the .creator attribute. We untangle it by first remembering the creator_id
# for all the streams and then removing those fields from the data.
# That allows us to successfully import streams, and then later after users
# are imported, we can set the .creator_id for all these streams correctly.
stream_id_to_creator_id = {}
for stream in data["zerver_stream"]:
creator_id = stream.pop("creator_id", None)
stream_id_to_creator_id[stream["id"]] = creator_id
if role_system_groups_dict is not None:
# Because the system user groups are missing, we manually set up
# the defaults for can_remove_subscribers_group for all the
@ -1136,6 +1147,13 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int = 1) -> Rea
user_profile.tos_version = UserProfile.TOS_VERSION_BEFORE_FIRST_LOGIN
UserProfile.objects.bulk_create(user_profiles)
# UserProfiles have been loaded, so now we're ready to set .creator_id
# for streams based on the mapping we saved earlier.
streams = Stream.objects.filter(id__in=stream_id_to_creator_id.keys())
for stream in streams:
stream.creator_id = stream_id_to_creator_id[stream.id]
Stream.objects.bulk_update(streams, ["creator_id"])
re_map_foreign_keys(data, "zerver_defaultstream", "stream", related_table="stream")
re_map_foreign_keys(data, "zerver_realmemoji", "author", related_table="user_profile")

View File

@ -763,6 +763,10 @@ class RealmImportExportTest(ExportFile):
cordelia = self.example_user("cordelia")
othello = self.example_user("othello")
denmark_stream = get_stream("Denmark", original_realm)
denmark_stream.creator = hamlet
denmark_stream.save()
internal_realm = get_realm(settings.SYSTEM_BOT_REALM)
cross_realm_bot = get_system_bot(settings.WELCOME_BOT, internal_realm.id)
@ -1012,6 +1016,12 @@ class RealmImportExportTest(ExportFile):
f'<p><span class="user-mention" data-user-id="{imported_polonius_user.id}">@Polonius</span></p>',
)
imported_hamlet_user = UserProfile.objects.get(
delivery_email=self.example_email("hamlet"), realm=imported_realm
)
imported_denmark_stream = Stream.objects.get(name="Denmark", realm=imported_realm)
self.assertEqual(imported_denmark_stream.creator, imported_hamlet_user)
# Check recipient_id was generated correctly for the imported users and streams.
for user_profile in UserProfile.objects.filter(realm=imported_realm):
self.assertEqual(