refactor(robot-server): Avoid features that will be removed in SQLAlchemy 2.0 (#16926)

Author: SyntaxColoring

robot-server/pytest.ini

@@ -4,3 +4,8 @@ markers =
         ot3_only: Test only functions using the OT3 hardware
 addopts = --color=yes --strict-markers
 asyncio_mode = auto
+
+# Don't allow any new code that uses features removed in SQLAlchemy 2.0.
+# We should remove this when we upgrade to SQLAlchemy 2.0.
+filterwarnings =
+    error::sqlalchemy.exc.RemovedIn20Warning

robot-server/robot_server/persistence/_migrations/_util.py

@@ -1,10 +1,11 @@
 """Shared helpers for migrations."""
 
-import sqlalchemy
-
 import shutil
+import typing
 from pathlib import Path
 
+import sqlalchemy
+
 from ..database import sqlite_rowid
 
 
@@ -54,3 +55,43 @@ def copytree_if_exists(src: Path, dst: Path) -> None:
         shutil.copytree(src=src, dst=dst)
     except FileNotFoundError:
         pass
+
+
+def add_column(
+    engine: sqlalchemy.engine.Engine,
+    table_name: str,
+    column: typing.Any,
+) -> None:
+    """Add a column to an existing SQL table, with an `ALTER TABLE` statement.
+
+    Params:
+        engine: A SQLAlchemy engine to connect to the database.
+        table_name: The SQL name of the parent table.
+        column: The SQLAlchemy column object.
+
+    Known limitations:
+
+    - This does not currently support indexes.
+    - This does not currently support constraints.
+    - The column will always be added as nullable. Adding non-nullable columns in
+      SQLite requires an elaborate and sensitive dance that we do not wish to attempt.
+      https://www.sqlite.org/lang_altertable.html#making_other_kinds_of_table_schema_changes
+
+    To avoid those limitations, instead of this function, consider this:
+
+    1. Start with an empty database, or drop or rename the current table.
+    2. Use SQLAlchemy's `metadata.create_all()` to create an empty table with the new
+       schema, including the new column.
+    3. Copy rows from the old table to the new one, populating the new column
+       however you please.
+    """
+    column_type = column.type.compile(engine.dialect)
+    with engine.begin() as transaction:
+        # todo(mm, 2024-11-25): This text seems like something that SQLAlchemy could generate for us
+        # (maybe: https://docs.sqlalchemy.org/en/20/core/metadata.html#sqlalchemy.schema.Column.compile),
+        # and that might help us account for indexes and constraints.
+        transaction.execute(
+            sqlalchemy.text(
+                f"ALTER TABLE {table_name} ADD COLUMN {column.key} {column_type}"
+            )
+        )

robot-server/robot_server/persistence/_migrations/v3_to_v4.py

@@ -9,10 +9,8 @@
 from pathlib import Path
 from contextlib import ExitStack
 import shutil
-from typing import Any
-
-import sqlalchemy
 
+from ._util import add_column
 from ..database import sql_engine_ctx
 from ..file_and_directory_names import DB_FILE
 from ..tables import schema_4
@@ -35,16 +33,6 @@ def migrate(self, source_dir: Path, dest_dir: Path) -> None:
             dest_engine = exit_stack.enter_context(sql_engine_ctx(dest_db_file))
             schema_4.metadata.create_all(dest_engine)
 
-            def add_column(
-                engine: sqlalchemy.engine.Engine,
-                table_name: str,
-                column: Any,
-            ) -> None:
-                column_type = column.type.compile(engine.dialect)
-                engine.execute(
-                    f"ALTER TABLE {table_name} ADD COLUMN {column.key} {column_type}"
-                )
-
             add_column(
                 dest_engine,
                 schema_4.analysis_table.name,

robot-server/robot_server/persistence/_migrations/v4_to_v5.py

@@ -9,10 +9,8 @@
 from pathlib import Path
 from contextlib import ExitStack
 import shutil
-from typing import Any
-
-import sqlalchemy
 
+from ._util import add_column
 from ..database import sql_engine_ctx
 from ..file_and_directory_names import DB_FILE
 from ..tables import schema_5
@@ -35,16 +33,6 @@ def migrate(self, source_dir: Path, dest_dir: Path) -> None:
             dest_engine = exit_stack.enter_context(sql_engine_ctx(dest_db_file))
             schema_5.metadata.create_all(dest_engine)
 
-            def add_column(
-                engine: sqlalchemy.engine.Engine,
-                table_name: str,
-                column: Any,
-            ) -> None:
-                column_type = column.type.compile(engine.dialect)
-                engine.execute(
-                    f"ALTER TABLE {table_name} ADD COLUMN {column.key} {column_type}"
-                )
-
             add_column(
                 dest_engine,
                 schema_5.protocol_table.name,

robot-server/robot_server/persistence/_migrations/v6_to_v7.py

@@ -11,10 +11,10 @@
 from pathlib import Path
 from contextlib import ExitStack
 import shutil
-from typing import Any
 
 import sqlalchemy
 
+from ._util import add_column
 from ..database import sql_engine_ctx, sqlite_rowid
 from ..tables import schema_7
 from .._folder_migrator import Migration
@@ -44,16 +44,6 @@ def migrate(self, source_dir: Path, dest_dir: Path) -> None:
 
             dest_transaction = exit_stack.enter_context(dest_engine.begin())
 
-            def add_column(
-                engine: sqlalchemy.engine.Engine,
-                table_name: str,
-                column: Any,
-            ) -> None:
-                column_type = column.type.compile(engine.dialect)
-                engine.execute(
-                    f"ALTER TABLE {table_name} ADD COLUMN {column.key} {column_type}"
-                )
-
             add_column(
                 dest_engine,
                 schema_7.run_command_table.name,

robot-server/robot_server/persistence/_migrations/v7_to_v8.py

@@ -10,10 +10,10 @@
 from pathlib import Path
 from contextlib import ExitStack
 import shutil
-from typing import Any
 
 import sqlalchemy
 
+from ._util import add_column
 from ..database import sql_engine_ctx
 from ..tables import schema_8
 from .._folder_migrator import Migration
@@ -41,16 +41,6 @@ def migrate(self, source_dir: Path, dest_dir: Path) -> None:
 
             dest_transaction = exit_stack.enter_context(dest_engine.begin())
 
-            def add_column(
-                engine: sqlalchemy.engine.Engine,
-                table_name: str,
-                column: Any,
-            ) -> None:
-                column_type = column.type.compile(engine.dialect)
-                engine.execute(
-                    f"ALTER TABLE {table_name} ADD COLUMN {column.key} {column_type}"
-                )
-
             add_column(
                 dest_engine,
                 schema_8.run_command_table.name,
@@ -73,9 +63,12 @@ def add_column(
 def _add_missing_indexes(dest_transaction: sqlalchemy.engine.Connection) -> None:
     # todo(2024-11-20): Probably add the indexes missing from prior migrations here.
     # https://opentrons.atlassian.net/browse/EXEC-827
-    dest_transaction.execute(
-        "CREATE UNIQUE INDEX ix_run_run_id_command_status_index_in_run ON run_command (run_id, command_status, index_in_run);"
+    index = next(
+        index
+        for index in schema_8.run_command_table.indexes
+        if index.name == "ix_run_run_id_command_status_index_in_run"
     )
+    index.create(dest_transaction)
 
 
 def _migrate_command_table_with_new_command_error_col_and_command_status(