From b4ce167a88d5bdc481e9e035483af4cdb170442b Mon Sep 17 00:00:00 2001 From: Mateusz Mandera Date: Sun, 15 Mar 2020 19:05:27 +0100 Subject: [PATCH] models: Add recipient foreign key to Huddle. This follows the already tested approach from 8acfa17fe629cf149dbdf4a59280e905bde82d49. --- analytics/tests/test_counts.py | 2 ++ zerver/lib/import_realm.py | 8 ++++++++ zerver/migrations/0270_huddle_recipient.py | 19 ++++++++++++++++++ ...0271_huddle_set_recipient_column_values.py | 20 +++++++++++++++++++ zerver/models.py | 12 +++++------ zerver/tests/test_import_export.py | 5 +++++ 6 files changed, 60 insertions(+), 6 deletions(-) create mode 100644 zerver/migrations/0270_huddle_recipient.py create mode 100644 zerver/migrations/0271_huddle_set_recipient_column_values.py diff --git a/analytics/tests/test_counts.py b/analytics/tests/test_counts.py index 2f8be5c973..117825c0e1 100644 --- a/analytics/tests/test_counts.py +++ b/analytics/tests/test_counts.py @@ -87,6 +87,8 @@ class AnalyticsTestCase(TestCase): kwargs[key] = kwargs.get(key, value) huddle = Huddle.objects.create(**kwargs) recipient = Recipient.objects.create(type_id=huddle.id, type=Recipient.HUDDLE) + huddle.recipient = recipient + huddle.save(update_fields=["recipient"]) return huddle, recipient def create_message(self, sender: UserProfile, recipient: Recipient, **kwargs: Any) -> Message: diff --git a/zerver/lib/import_realm.py b/zerver/lib/import_realm.py index 34bcd29784..ea22b6a316 100644 --- a/zerver/lib/import_realm.py +++ b/zerver/lib/import_realm.py @@ -513,6 +513,10 @@ def remove_denormalized_recipient_column_from_data(data: TableData) -> None: if 'recipient' in user_profile_dict: del user_profile_dict['recipient'] + for huddle_dict in data['zerver_huddle']: + if 'recipient' in huddle_dict: + del huddle_dict['recipient'] + def get_db_table(model_class: Any) -> str: """E.g. (RealmDomain -> 'zerver_realmdomain')""" return model_class._meta.db_table @@ -914,6 +918,10 @@ def do_import_realm(import_dir: Path, subdomain: str, processes: int=1) -> Realm if 'zerver_huddle' in data: process_huddle_hash(data, 'zerver_huddle') bulk_import_model(data, Huddle) + for huddle in Huddle.objects.filter(recipient_id=None): + recipient = Recipient.objects.get(type=Recipient.HUDDLE, type_id=huddle.id) + huddle.recipient = recipient + huddle.save(update_fields=["recipient"]) if 'zerver_userhotspot' in data: fix_datetime_fields(data, 'zerver_userhotspot') diff --git a/zerver/migrations/0270_huddle_recipient.py b/zerver/migrations/0270_huddle_recipient.py new file mode 100644 index 0000000000..9eac7771d7 --- /dev/null +++ b/zerver/migrations/0270_huddle_recipient.py @@ -0,0 +1,19 @@ +# Generated by Django 2.2.10 on 2020-03-15 17:25 + +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + + dependencies = [ + ('zerver', '0269_gitlab_auth'), + ] + + operations = [ + migrations.AddField( + model_name='huddle', + name='recipient', + field=models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, to='zerver.Recipient'), + ), + ] diff --git a/zerver/migrations/0271_huddle_set_recipient_column_values.py b/zerver/migrations/0271_huddle_set_recipient_column_values.py new file mode 100644 index 0000000000..493da42932 --- /dev/null +++ b/zerver/migrations/0271_huddle_set_recipient_column_values.py @@ -0,0 +1,20 @@ +# -*- coding: utf-8 -*- +from django.db import migrations + + +class Migration(migrations.Migration): + + dependencies = [ + ('zerver', '0270_huddle_recipient'), + ] + + operations = [ + migrations.RunSQL( + """ + UPDATE zerver_huddle + SET recipient_id = zerver_recipient.id + FROM zerver_recipient + WHERE zerver_recipient.type_id = zerver_huddle.id AND zerver_recipient.type = 3; + """, + reverse_sql='UPDATE zerver_huddle SET recipient_id = NULL'), + ] diff --git a/zerver/models.py b/zerver/models.py index 4ddec17242..2f1d6af735 100644 --- a/zerver/models.py +++ b/zerver/models.py @@ -1547,13 +1547,9 @@ def bulk_get_streams(realm: Realm, stream_names: STREAM_NAMES) -> Dict[str, Any] [stream_name.lower() for stream_name in stream_names], id_fetcher=stream_to_lower_name) -def get_recipient_cache_key(type: int, type_id: int) -> str: +def get_recipient_cache_key(type: int, type_id: int) -> str: # nocoverage return u"%s:get_recipient:%s:%s" % (cache.KEY_PREFIX, type, type_id,) -@cache_with_key(get_recipient_cache_key, timeout=3600*24*7) -def get_recipient(type: int, type_id: int) -> Recipient: - return Recipient.objects.get(type_id=type_id, type=type) - def get_huddle_recipient(user_profile_ids: Set[int]) -> Recipient: # The caller should ensure that user_profile_ids includes @@ -1561,7 +1557,7 @@ def get_huddle_recipient(user_profile_ids: Set[int]) -> Recipient: # we hit another cache to get the recipient. We may want to # unify our caching strategy here. huddle = get_huddle(list(user_profile_ids)) - return get_recipient(Recipient.HUDDLE, huddle.id) + return huddle.recipient def get_huddle_user_ids(recipient: Recipient) -> List[int]: assert(recipient.type == Recipient.HUDDLE) @@ -2287,6 +2283,8 @@ class Huddle(models.Model): # TODO: We should consider whether using # CommaSeparatedIntegerField would be better. huddle_hash = models.CharField(max_length=40, db_index=True, unique=True) # type: str + # Foreign key to the Recipient object for this Huddle. + recipient = models.ForeignKey(Recipient, null=True, on_delete=models.SET_NULL) def get_huddle_hash(id_list: List[int]) -> str: id_list = sorted(set(id_list)) @@ -2307,6 +2305,8 @@ def get_huddle_backend(huddle_hash: str, id_list: List[int]) -> Huddle: if created: recipient = Recipient.objects.create(type_id=huddle.id, type=Recipient.HUDDLE) + huddle.recipient = recipient + huddle.save(update_fields=["recipient"]) subs_to_create = [Subscription(recipient=recipient, user_profile_id=user_profile_id) for user_profile_id in id_list] diff --git a/zerver/tests/test_import_export.py b/zerver/tests/test_import_export.py index 6bcc676157..40a6f6bedf 100644 --- a/zerver/tests/test_import_export.py +++ b/zerver/tests/test_import_export.py @@ -1013,6 +1013,11 @@ class ImportExportTest(ZulipTestCase): self.assertEqual(stream.recipient_id, Recipient.objects.get(type=Recipient.STREAM, type_id=stream.id).id) + for huddle_object in Huddle.objects.all(): + # Huddles don't have a realm column, so we just test all Huddles for simplicity. + self.assertEqual(huddle_object.recipient_id, Recipient.objects.get(type=Recipient.HUDDLE, + type_id=huddle_object.id).id) + def test_import_files_from_local(self) -> None: realm = Realm.objects.get(string_id='zulip')