diff --git a/zerver/lib/test_classes.py b/zerver/lib/test_classes.py index 1aa337ab96..13c800ccf3 100644 --- a/zerver/lib/test_classes.py +++ b/zerver/lib/test_classes.py @@ -2277,7 +2277,7 @@ one or more new messages. return body -class MigrationsTestCase(ZulipTestCase): # nocoverage +class MigrationsTestCase(ZulipTransactionTestCase): # nocoverage """ Test class for database migrations inspired by this blog post: https://www.caktusgroup.com/blog/2016/02/02/writing-unit-tests-django-migrations/ @@ -2295,6 +2295,7 @@ class MigrationsTestCase(ZulipTestCase): # nocoverage @override def setUp(self) -> None: + super().setUp() assert ( self.migrate_from and self.migrate_to ), f"TestCase '{type(self).__name__}' must define migrate_from and migrate_to properties" diff --git a/zerver/migrations/0485_alter_usermessage_flags_and_add_index.py b/zerver/migrations/0485_alter_usermessage_flags_and_add_index.py new file mode 100644 index 0000000000..92032de244 --- /dev/null +++ b/zerver/migrations/0485_alter_usermessage_flags_and_add_index.py @@ -0,0 +1,69 @@ +# Generated by Django 4.2.6 on 2023-10-19 06:37 + +import bitfield.models +from django.contrib.postgres.operations import AddIndexConcurrently +from django.db import migrations, models + + +class Migration(migrations.Migration): + atomic = False + + dependencies = [ + ("zerver", "0484_preregistrationrealm_default_language"), + ] + + operations = [ + migrations.AlterField( + model_name="archivedusermessage", + name="flags", + field=bitfield.models.BitField( + [ + "read", + "starred", + "collapsed", + "mentioned", + "wildcard_mentioned", + "topic_wildcard_mentioned", + "group_mentioned", + "force_expand", + "force_collapse", + "has_alert_word", + "historical", + "is_private", + "active_mobile_push_notification", + ], + default=0, + ), + ), + migrations.AlterField( + model_name="usermessage", + name="flags", + field=bitfield.models.BitField( + [ + "read", + "starred", + "collapsed", + "mentioned", + "wildcard_mentioned", + "topic_wildcard_mentioned", + "group_mentioned", + "force_expand", + "force_collapse", + "has_alert_word", + "historical", + "is_private", + "active_mobile_push_notification", + ], + default=0, + ), + ), + AddIndexConcurrently( + model_name="usermessage", + index=models.Index( + "user_profile", + "message", + condition=models.Q(("flags__andnz", 120)), + name="zerver_usermessage_any_mentioned_message_id", + ), + ), + ] diff --git a/zerver/models.py b/zerver/models.py index 796264ac32..f971d7c93a 100644 --- a/zerver/models.py +++ b/zerver/models.py @@ -3476,9 +3476,9 @@ class AbstractUserMessage(models.Model): "collapsed", "mentioned", "wildcard_mentioned", - # These next 4 flags are from features that have since been removed. - "summarize_in_home", - "summarize_in_stream", + "topic_wildcard_mentioned", + "group_mentioned", + # These next 2 flags are from features that have since been removed. "force_expand", "force_collapse", # Whether the message contains any of the user's alert words. @@ -3506,12 +3506,12 @@ class AbstractUserMessage(models.Model): "has_alert_word", "mentioned", "wildcard_mentioned", + "topic_wildcard_mentioned", + "group_mentioned", "historical", # Unused flags can't be edited. "force_expand", "force_collapse", - "summarize_in_home", - "summarize_in_stream", } flags: BitHandler = BitField(flags=ALL_FLAGS, default=0) @@ -3607,6 +3607,17 @@ class UserMessage(AbstractUserMessage): | Q(flags__andnz=AbstractUserMessage.flags.wildcard_mentioned.mask), name="zerver_usermessage_wildcard_mentioned_message_id", ), + models.Index( + "user_profile", + "message", + condition=Q( + flags__andnz=AbstractUserMessage.flags.mentioned.mask + | AbstractUserMessage.flags.wildcard_mentioned.mask + | AbstractUserMessage.flags.topic_wildcard_mentioned.mask + | AbstractUserMessage.flags.group_mentioned.mask + ), + name="zerver_usermessage_any_mentioned_message_id", + ), models.Index( "user_profile", "message", diff --git a/zerver/tests/test_migrations.py b/zerver/tests/test_migrations.py index a71b263414..ffe440382e 100644 --- a/zerver/tests/test_migrations.py +++ b/zerver/tests/test_migrations.py @@ -4,10 +4,7 @@ # You can also read # https://www.caktusgroup.com/blog/2016/02/02/writing-unit-tests-django-migrations/ # to get a tutorial on the framework that inspired this feature. -from datetime import datetime, timezone -from unittest import skip - -import orjson +from django.db import connections from django.db.migrations.state import StateApps from typing_extensions import override @@ -30,77 +27,26 @@ from zerver.lib.test_helpers import use_db_models # been tested for a migration being merged. -@skip("Cannot be run because there is a non-atomic migration that has been merged after it") -class ScheduledEmailData(MigrationsTestCase): - migrate_from = "0467_rename_extradata_realmauditlog_extra_data_json" - migrate_to = "0468_rename_followup_day_email_templates" +class UserMessageIndex(MigrationsTestCase): + migrate_from = "0484_preregistrationrealm_default_language" + migrate_to = "0485_alter_usermessage_flags_and_add_index" + + def index_exists(self, index_name: str) -> bool: + table_name = "zerver_usermessage" + connection = connections["default"] + + with connection.cursor() as cursor: + # We use parameterized query to prevent SQL injection vulnerabilities + query = "SELECT indexname FROM pg_indexes WHERE tablename = %s AND indexname = %s" + cursor.execute(query, [table_name, index_name]) + return cursor.fetchone() is not None @use_db_models @override def setUpBeforeMigration(self, apps: StateApps) -> None: - iago = self.example_user("iago") - ScheduledEmail = apps.get_model("zerver", "ScheduledEmail") - send_date = datetime(2025, 1, 1, 1, 0, 0, 0, tzinfo=timezone.utc) + self.assertTrue(self.index_exists("zerver_usermessage_wildcard_mentioned_message_id")) + self.assertFalse(self.index_exists("zerver_usermessage_any_mentioned_message_id")) - templates = [ - ["zerver/emails/followup_day1", "a", True, 10], - ["zerver/emails/followup_day2", "b", False, 20], - ["zerver/emails/onboarding_zulip_guide", "c", True, 30], - ] - - for template in templates: - email_fields = { - "template_prefix": template[0], - "string_context": template[1], - "boolean_context": template[2], - "integer_context": template[3], - } - - email = ScheduledEmail.objects.create( - type=1, - realm=iago.realm, - scheduled_timestamp=send_date, - data=orjson.dumps(email_fields).decode(), - ) - email.users.add(iago.id) - - def test_updated_email_templates(self) -> None: - ScheduledEmail = self.apps.get_model("zerver", "ScheduledEmail") - send_date = datetime(2025, 1, 1, 1, 0, 0, 0, tzinfo=timezone.utc) - - old_templates = [ - "zerver/emails/followup_day1", - "zerver/emails/followup_day2", - ] - - current_templates = [ - "zerver/emails/account_registered", - "zerver/emails/onboarding_zulip_guide", - "zerver/emails/onboarding_zulip_topics", - ] - - email_data = [ - ["zerver/emails/account_registered", "a", True, 10], - ["zerver/emails/onboarding_zulip_topics", "b", False, 20], - ["zerver/emails/onboarding_zulip_guide", "c", True, 30], - ] - - scheduled_emails = ScheduledEmail.objects.all() - self.assert_length(scheduled_emails, 3) - - checked_emails = [] - for email in scheduled_emails: - self.assertEqual(email.type, 1) - self.assertEqual(email.scheduled_timestamp, send_date) - - updated_data = orjson.loads(email.data) - template_prefix = updated_data["template_prefix"] - self.assertFalse(template_prefix in old_templates) - for data in email_data: - if template_prefix == data[0]: - self.assertEqual(updated_data["string_context"], data[1]) - self.assertEqual(updated_data["boolean_context"], data[2]) - self.assertEqual(updated_data["integer_context"], data[3]) - checked_emails.append(template_prefix) - - self.assertEqual(current_templates, sorted(checked_emails)) + def test_new_index_created_old_index_not_removed(self) -> None: + self.assertTrue(self.index_exists("zerver_usermessage_wildcard_mentioned_message_id")) + self.assertTrue(self.index_exists("zerver_usermessage_any_mentioned_message_id"))