Skip to content

Commit

Permalink
import: Fix import of Stream.creator field.
Browse files Browse the repository at this point in the history
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".
```
  • Loading branch information
mateuszmandera authored and timabbott committed May 14, 2024
1 parent 2e8fdab commit 160076f
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 0 deletions.
18 changes: 18 additions & 0 deletions zerver/lib/import_realm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")

Expand Down
10 changes: 10 additions & 0 deletions zerver/tests/test_import_export.py
Original file line number Diff line number Diff line change
Expand Up @@ -763,6 +763,10 @@ def test_import_realm(self) -> None:
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)

Expand Down Expand Up @@ -1012,6 +1016,12 @@ def assert_realm_values(f: Callable[[Realm], object]) -> None:
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(
Expand Down

0 comments on commit 160076f

Please sign in to comment.