diff --git a/robot-server/robot_server/persistence/_database.py b/robot-server/robot_server/persistence/_database.py index f38d4f22f8e..3055898dd36 100644 --- a/robot-server/robot_server/persistence/_database.py +++ b/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,35 +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:/// - # where 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 diff --git a/server-utils/server_utils/sql_utils.py b/server-utils/server_utils/sql_utils.py new file mode 100644 index 00000000000..4033ead0c5b --- /dev/null +++ b/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:/// + # where 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") diff --git a/system-server/system_server/persistence/database.py b/system-server/system_server/persistence/database.py index 28de145489f..2f74cd6da7e 100644 --- a/system-server/system_server/persistence/database.py +++ b/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:/// - # where 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