Deduplicate with system-server.

robot-server/robot_server/persistence/_database.py

@@ -3,6 +3,8 @@
 
 import sqlalchemy
 
+from server_utils import sql_utils
+
 from ._tables import add_tables_to_db
 from ._migrations import migrate
 
@@ -29,46 +31,16 @@ def create_sql_engine(path: Path) -> sqlalchemy.engine.Engine:
         Migrations can take several minutes. If calling this from an async function,
         offload this to a thread to avoid blocking the event loop.
     """
-    sql_engine = _open_db_no_cleanup(db_file_path=path)
+    sql_engine = sqlalchemy.create_engine(sql_utils.get_connection_url(path))
 
     try:
+        sql_utils.enable_foreign_key_constraints(sql_engine)
+        sql_utils.fix_transactions(sql_engine)
         add_tables_to_db(sql_engine)
         migrate(sql_engine)
+
     except Exception:
         sql_engine.dispose()
         raise
 
     return sql_engine
-
-
-def _open_db_no_cleanup(db_file_path: Path) -> sqlalchemy.engine.Engine:
-    """Create a database engine for performing transactions."""
-    engine = sqlalchemy.create_engine(
-        # sqlite://<hostname>/<path>
-        # where <hostname> is empty.
-        f"sqlite:///{db_file_path}",
-    )
-
-    # Enable foreign key support in sqlite
-    # https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#foreign-key-support
-    @sqlalchemy.event.listens_for(engine, "connect")  # type: ignore[misc]
-    def _set_sqlite_pragma(
-        dbapi_connection: sqlalchemy.engine.CursorResult,
-        connection_record: sqlalchemy.engine.CursorResult,
-    ) -> None:
-        cursor = dbapi_connection.cursor()
-        cursor.execute("PRAGMA foreign_keys=ON;")
-        cursor.close()
-
-    @sqlalchemy.event.listens_for(engine, "connect")  # type: ignore[misc]
-    def do_connect(dbapi_connection, connection_record):
-        # disable pysqlite's emitting of the BEGIN statement entirely.
-        # also stops it from emitting COMMIT before any DDL.
-        dbapi_connection.isolation_level = None
-
-    @sqlalchemy.event.listens_for(engine, "begin")  # type: ignore[misc]
-    def do_begin(conn):
-        # emit our own BEGIN
-        conn.exec_driver_sql("BEGIN")
-
-    return engine

server-utils/server_utils/sql_utils.py

@@ -0,0 +1,76 @@
+"""Utilities for working with SQLite databases through SQLAlchemy."""
+
+from pathlib import Path
+from typing import Any
+
+import sqlalchemy
+
+
+def get_connection_url(db_file_path: Path) -> str:
+    """Return a connection URL to pass to `sqlalchemy.create_engine()`.
+
+    Params:
+        db_file_path: The path to the SQLite database file to open.
+        (This file often has an extension like .db, .sqlite, or .sqlite3.)
+    """
+    # sqlite://<hostname>/<path>
+    # where <hostname> is empty.
+    return f"sqlite:///{db_file_path}"
+
+
+def enable_foreign_key_constraints(engine: sqlalchemy.engine.Engine) -> None:
+    """Enable SQLite's enforcement of foreign key constraints.
+
+    SQLite does not enforce foreign key constraints by default, for backwards compatibility.
+
+    This should be called once per SQLAlchemy engine, shortly after creating it,
+    before doing anything substantial with it.
+
+    Params:
+        engine: A SQLAlchemy engine connected to a SQLite database.
+    """
+    # Copied from:
+    # https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#foreign-key-support
+
+    @sqlalchemy.event.listens_for(engine, "connect")  # type: ignore[misc]
+    def on_connect(
+        # TODO(mm, 2023-08-29): Improve these type annotations when we have SQLAlchemy 2.0.
+        dbapi_connection: Any,
+        connection_record: object,
+    ) -> None:
+        cursor = dbapi_connection.cursor()
+        cursor.execute("PRAGMA foreign_keys=ON;")
+        cursor.close()
+
+
+def fix_transactions(engine: sqlalchemy.engine.Engine) -> None:
+    """Make SQLite transactions behave sanely.
+
+    This works around various misbehaviors in Python's `sqlite3` driver (aka `pysqlite`),
+    which is a middle layer between SQLAlchemy and the underlying SQLite library.
+    These misbehaviors can make transactions not actually behave transactionally. See:
+    https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#serializable-isolation-savepoints-transactional-ddl
+
+    This should be called once per SQLAlchemy engine, shortly after creating it,
+    before doing anything substantial with it.
+
+    Params:
+        engine: A SQLAlchemy engine connected to a SQLite database.
+    """
+    # Copied from:
+    # https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#serializable-isolation-savepoints-transactional-ddl.
+
+    @sqlalchemy.event.listens_for(engine, "connect")  # type: ignore[misc]
+    def on_connect(
+        # TODO(mm, 2023-08-29): Improve these type annotations when we have SQLAlchemy 2.0.
+        dbapi_connection: Any,
+        connection_record: object,
+    ) -> None:
+        # disable pysqlite's emitting of the BEGIN statement entirely.
+        # also stops it from emitting COMMIT before any DDL.
+        dbapi_connection.isolation_level = None
+
+    @sqlalchemy.event.listens_for(engine, "begin")  # type: ignore[misc]
+    def on_begin(conn: sqlalchemy.engine.Connection) -> None:
+        # emit our own BEGIN
+        conn.exec_driver_sql("BEGIN")

system-server/system_server/persistence/database.py

@@ -2,6 +2,9 @@
 from pathlib import Path
 
 import sqlalchemy
+
+from server_utils import sql_utils
+
 from .tables import add_tables_to_db
 from .migrations import migrate
 
@@ -23,35 +26,16 @@
 
 def create_sql_engine(path: Path) -> sqlalchemy.engine.Engine:
     """Create a SQL engine with tables and migrations."""
-    sql_engine = _open_db(db_file_path=path)
+    sql_engine = sqlalchemy.create_engine(sql_utils.get_connection_url(path))
 
     try:
+        sql_utils.enable_foreign_key_constraints(sql_engine)
+        sql_utils.fix_transactions(sql_engine)
         add_tables_to_db(sql_engine)
         migrate(sql_engine)
+
     except Exception:
         sql_engine.dispose()
         raise
 
     return sql_engine
-
-
-def _open_db(db_file_path: Path) -> sqlalchemy.engine.Engine:
-    """Create a database engine for performing transactions."""
-    engine = sqlalchemy.create_engine(
-        # sqlite://<hostname>/<path>
-        # where <hostname> is empty.
-        f"sqlite:///{db_file_path}",
-    )
-
-    # Enable foreign key support in sqlite
-    # https://docs.sqlalchemy.org/en/14/dialects/sqlite.html#foreign-key-support
-    @sqlalchemy.event.listens_for(engine, "connect")  # type: ignore[misc]
-    def _set_sqlite_pragma(
-        dbapi_connection: sqlalchemy.engine.CursorResult,
-        connection_record: sqlalchemy.engine.CursorResult,
-    ) -> None:
-        cursor = dbapi_connection.cursor()
-        cursor.execute("PRAGMA foreign_keys=ON;")
-        cursor.close()
-
-    return engine