Automatically port sqlite indexes

Author: erikjohnston

synapse/storage/background_updates.py

@@ -14,7 +14,6 @@
 import abc
 import logging
 from enum import Enum, IntEnum
-from io import StringIO
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
@@ -1055,15 +1054,15 @@ def run_validate_constraint_and_delete_rows_schema_delta(
     constraint: Constraint,
     sqlite_table_name: str,
     sqlite_table_schema: str,
-    sqlite_post_schema: Optional[str],
 ) -> None:
     """Runs a schema delta to add a constraint to the table. This should be run
     in a schema delta file.
 
     For PostgreSQL the constraint is added and validated in the background.
 
     For SQLite the table is recreated and data copied across immediately. This
-    is done by the caller passing in a script to create the new table.
+    is done by the caller passing in a script to create the new table. Note that
+    table indexes and triggers are copied over automatically.
 
     There must be a corresponding call to
     `register_background_validate_constraint_and_delete_rows` to register the
@@ -1075,8 +1074,6 @@ def run_validate_constraint_and_delete_rows_schema_delta(
         new constraint constraint: A `Constraint` object describing the
         constraint sqlite_table_name: For SQLite the name of the empty copy of
         table sqlite_table_schema: A SQL script for creating the above table.
-        sqlite_post_schema: A SQL script run after migration, to add back
-          indices and the like.
     """
 
     if isinstance(txn.database_engine, PostgresEngine):
@@ -1097,14 +1094,25 @@ def run_validate_constraint_and_delete_rows_schema_delta(
         )
     else:
         # For SQLite, we:
-        #   1. create an empty copy of the table
-        #   2. copy across the rows (that satisfy the check)
-        #   3. replace the old table with the new able.
-
-        # We import this here to avoid circular imports.
-        from synapse.storage.prepare_database import execute_statements_from_stream
+        #   1. fetch all indexes/triggers/etc related to the table
+        #   2. create an empty copy of the table
+        #   3. copy across the rows (that satisfy the check)
+        #   4. replace the old table with the new able.
+        #   5. add back all the indexes/triggers/etc
+
+        # Fetch the indexes/triggers/etc. Note that `sql` column being null is
+        # due to indexes being auto created based on the class definition (e.g.
+        # PRIMARY KEY), and so don't need to be recreated.
+        txn.execute(
+            """
+            SELECT sql FROM sqlite_schema
+            WHERE tbl_name = ? AND type != 'table' AND sql IS NOT NULL
+            """,
+            (table,),
+        )
+        extras = [row[0] for row in txn]
 
-        execute_statements_from_stream(txn, StringIO(sqlite_table_schema))
+        txn.execute(sqlite_table_schema)
 
         sql = f"""
             INSERT INTO {sqlite_table_name} SELECT * FROM {table}
@@ -1116,5 +1124,5 @@ def run_validate_constraint_and_delete_rows_schema_delta(
         txn.execute(f"DROP TABLE {table}")
         txn.execute(f"ALTER TABLE {sqlite_table_name} RENAME TO {table}")
 
-        if sqlite_post_schema:
-            execute_statements_from_stream(txn, StringIO(sqlite_post_schema))
+        for extra in extras:
+            txn.execute(extra)

synapse/storage/schema/main/delta/78/03event_extremities_constraints.py

@@ -29,12 +29,7 @@
         room_id TEXT NOT NULL,
         UNIQUE (event_id, room_id),
         CONSTRAINT event_forward_extremities_event_id FOREIGN KEY (event_id) REFERENCES events (event_id)
-    );
-"""
-
-FORWARD_EXTREMITIES_INDICES_SCHEMA = """
-    CREATE INDEX ev_extrem_room ON event_forward_extremities(room_id);
-    CREATE INDEX ev_extrem_id ON event_forward_extremities(event_id);
+    )
 """
 
 
@@ -48,7 +43,6 @@ def run_create(cur: LoggingTransaction, database_engine: BaseDatabaseEngine) ->
         constraint=ForeignKeyConstraint("events", [("event_id", "event_id")]),
         sqlite_table_name="event_forward_extremities2",
         sqlite_table_schema=FORWARD_EXTREMITIES_TABLE_SCHEMA,
-        sqlite_post_schema=FORWARD_EXTREMITIES_INDICES_SCHEMA,
     )
 
     # We can't add a similar constraint to `event_backward_extremities` as the

tests/storage/test_background_update.py

@@ -27,7 +27,7 @@
     run_validate_constraint_and_delete_rows_schema_delta,
 )
 from synapse.storage.database import LoggingTransaction
-from synapse.storage.engines import PostgresEngine
+from synapse.storage.engines import PostgresEngine, Sqlite3Engine
 from synapse.types import JsonDict
 from synapse.util import Clock
 
@@ -440,6 +440,15 @@ def test_not_null_constraint(self) -> None:
             )
         )
 
+        # We add an index so that we can check that its correctly recreated when
+        # using SQLite.
+        index_sql = "CREATE INDEX test_index ON test_constraint(a)"
+        self.get_success(
+            self.store.db_pool.execute(
+                "test_not_null_constraint", lambda _: None, index_sql
+            )
+        )
+
         self.get_success(
             self.store.db_pool.simple_insert("test_constraint", {"a": 1, "b": 1})
         )
@@ -470,7 +479,6 @@ def delta(txn: LoggingTransaction) -> None:
                 constraint=NotNullConstraint("b"),
                 sqlite_table_name="test_constraint2",
                 sqlite_table_schema=table2_sqlite,
-                sqlite_post_schema=None,
             )
 
         self.get_success(
@@ -513,6 +521,17 @@ def delta(txn: LoggingTransaction) -> None:
             exc=self.store.database_engine.module.IntegrityError,
         )
 
+        # Check the index is still there for SQLite.
+        if isinstance(self.store.database_engine, Sqlite3Engine):
+            # Ensure the index exists in the schema.
+            self.get_success(
+                self.store.db_pool.simple_select_one_onecol(
+                    table="sqlite_schema",
+                    keyvalues={"tbl_name": "test_constraint"},
+                    retcol="name",
+                )
+            )
+
     def test_foreign_constraint(self) -> None:
         """Tests adding a not foreign key constraint."""
 
@@ -570,7 +589,6 @@ def delta(txn: LoggingTransaction) -> None:
                 constraint=ForeignKeyConstraint("base_table", [("b", "b")]),
                 sqlite_table_name="test_constraint2",
                 sqlite_table_schema=table2_sqlite,
-                sqlite_post_schema=None,
             )
 
         self.get_success(