models: Manage indexes from migration 0001 with Django.

Signed-off-by: Anders Kaseorg <anders@zulip.com>

zerver/lib/retention.py

@@ -89,6 +89,8 @@ models_with_message_key: List[Dict[str, Any]] = [
     },
 ]
 
+EXCLUDE_FIELDS = {Message._meta.get_field("search_tsvector")}
+
 
 @transaction.atomic(savepoint=False)
 def move_rows(
@@ -104,8 +106,9 @@ def move_rows(
         # Use base_model's db_table unless otherwise specified.
         src_db_table = base_model._meta.db_table
 
-    src_fields = [Identifier(src_db_table, field.column) for field in base_model._meta.fields]
-    dst_fields = [Identifier(field.column) for field in base_model._meta.fields]
+    fields = [field for field in base_model._meta.fields if field not in EXCLUDE_FIELDS]
+    src_fields = [Identifier(src_db_table, field.column) for field in fields]
+    dst_fields = [Identifier(field.column) for field in fields]
     sql_args = {
         "src_fields": SQL(",").join(src_fields),
         "dst_fields": SQL(",").join(dst_fields),

zerver/migrations/0001_initial.py

@@ -5,9 +5,12 @@ import django.core.validators
 import django.db.models.deletion
 import django.utils.timezone
 from django.conf import settings
+from django.contrib.postgres.indexes import GinIndex
+from django.contrib.postgres.search import SearchVectorField
 from django.db import migrations, models
 from django.db.backends.postgresql.schema import DatabaseSchemaEditor
 from django.db.migrations.state import StateApps
+from django.db.models.functions import Upper
 
 from zerver.models import generate_email_token_for_stream
 
@@ -62,10 +65,6 @@ CREATE FUNCTION escape_html(text) RETURNS text IMMUTABLE LANGUAGE 'sql' AS $$
                                  '>', '>'), '"', '"'), '''', ''');
 $$ ;
 
-ALTER TABLE zerver_message ADD COLUMN search_tsvector tsvector;
-CREATE INDEX zerver_message_search_tsvector ON zerver_message USING gin(search_tsvector);
-ALTER INDEX zerver_message_search_tsvector SET (fastupdate = OFF);
-
 CREATE TABLE fts_update_log (id SERIAL PRIMARY KEY, message_id INTEGER NOT NULL);
 CREATE FUNCTION do_notify_fts_update_log() RETURNS trigger LANGUAGE plpgsql AS
   $$ BEGIN NOTIFY fts_update_log; RETURN NEW; END $$;
@@ -706,6 +705,17 @@ CREATE TRIGGER zerver_message_update_search_tsvector_async
                 verbose_name="user permissions",
             ),
         ),
+        migrations.AddField(
+            model_name="message",
+            name="search_tsvector",
+            field=SearchVectorField(null=True),
+        ),
+        migrations.AddIndex(
+            model_name="message",
+            index=GinIndex(
+                "search_tsvector", fastupdate=False, name="zerver_message_search_tsvector"
+            ),
+        ),
         migrations.RunSQL(
             sql=fts_sql,
         ),
@@ -745,13 +755,13 @@ CREATE TRIGGER zerver_message_update_search_tsvector_async
             name="last_login",
             field=models.DateTimeField(blank=True, null=True, verbose_name="last login"),
         ),
-        migrations.RunSQL(
-            sql="CREATE INDEX upper_subject_idx ON zerver_message ((upper(subject)));",
-            reverse_sql="DROP INDEX upper_subject_idx;",
+        migrations.AddIndex(
+            model_name="message",
+            index=models.Index(Upper("subject"), name="upper_subject_idx"),
         ),
-        migrations.RunSQL(
-            sql="CREATE INDEX upper_stream_name_idx ON zerver_stream ((upper(name)));",
-            reverse_sql="DROP INDEX upper_stream_name_idx;",
+        migrations.AddIndex(
+            model_name="stream",
+            index=models.Index(Upper("name"), name="upper_stream_name_idx"),
         ),
         migrations.AddField(
             model_name="userprofile",
@@ -767,9 +777,9 @@ CREATE TRIGGER zerver_message_update_search_tsvector_async
                 )
             },
         ),
-        migrations.RunSQL(
-            sql="CREATE INDEX upper_userprofile_email_idx ON zerver_userprofile ((upper(email)));",
-            reverse_sql="DROP INDEX upper_userprofile_email_idx;",
+        migrations.AddIndex(
+            model_name="userprofile",
+            index=models.Index(Upper("email"), name="upper_userprofile_email_idx"),
         ),
         migrations.AlterField(
             model_name="userprofile",
@@ -781,9 +791,9 @@ CREATE TRIGGER zerver_message_update_search_tsvector_async
             name="is_bot",
             field=models.BooleanField(db_index=True, default=False),
         ),
-        migrations.RunSQL(
-            sql="CREATE INDEX upper_preregistration_email_idx ON zerver_preregistrationuser ((upper(email)));",
-            reverse_sql="DROP INDEX upper_preregistration_email_idx;",
+        migrations.AddIndex(
+            model_name="preregistrationuser",
+            index=models.Index(Upper("email"), name="upper_preregistration_email_idx"),
         ),
         migrations.AlterField(
             model_name="userprofile",

zerver/models.py

@@ -26,6 +26,8 @@ from bitfield.types import BitHandler
 from django.conf import settings
 from django.contrib.auth.models import AbstractBaseUser, PermissionsMixin, UserManager
 from django.contrib.contenttypes.fields import GenericRelation
+from django.contrib.postgres.indexes import GinIndex
+from django.contrib.postgres.search import SearchVectorField
 from django.core.exceptions import ValidationError
 from django.core.validators import MinLengthValidator, RegexValidator, URLValidator, validate_email
 from django.db import models, transaction
@@ -2061,6 +2063,11 @@ class UserProfile(AbstractBaseUser, PermissionsMixin, UserBaseSettings):
 
         super().set_password(password)
 
+    class Meta:
+        indexes = [
+            models.Index(Upper("email"), name="upper_userprofile_email_idx"),
+        ]
+
 
 class PasswordTooWeakError(Exception):
     pass
@@ -2169,6 +2176,11 @@ class PreregistrationUser(models.Model):
     )
     invited_as: int = models.PositiveSmallIntegerField(default=INVITE_AS["MEMBER"])
 
+    class Meta:
+        indexes = [
+            models.Index(Upper("email"), name="upper_preregistration_email_idx"),
+        ]
+
 
 def filter_to_valid_prereg_users(
     query: QuerySet,
@@ -2415,6 +2427,11 @@ class Stream(models.Model):
         result["is_announcement_only"] = self.stream_post_policy == Stream.STREAM_POST_POLICY_ADMINS
         return result
 
+    class Meta:
+        indexes = [
+            models.Index(Upper("name"), name="upper_stream_name_idx"),
+        ]
+
 
 post_save.connect(flush_stream, sender=Stream)
 post_delete.connect(flush_stream, sender=Stream)
@@ -2724,6 +2741,7 @@ class ArchivedMessage(AbstractMessage):
 
 class Message(AbstractMessage):
     id: int = models.AutoField(auto_created=True, primary_key=True, verbose_name="ID")
+    search_tsvector = SearchVectorField(null=True)
 
     def topic_name(self) -> str:
         """
@@ -2802,6 +2820,8 @@ class Message(AbstractMessage):
 
     class Meta:
         indexes = [
+            GinIndex("search_tsvector", fastupdate=False, name="zerver_message_search_tsvector"),
+            models.Index(Upper("subject"), name="upper_subject_idx"),
             models.Index("date_sent", name="zerver_message_date_sent_3b5b05d8"),
             models.Index(
                 "recipient",