Skip to content

Commit

Permalink
fix(robot-server,system-server): Make SQL transactions behave sanely (#…
Browse files Browse the repository at this point in the history
…13424)

* Use SQLAlchemy's workarounds to fix transactions.

* Deduplicate with system-server.
  • Loading branch information
SyntaxColoring authored and TamarZanzouri committed Sep 13, 2023
1 parent 0331a21 commit 5f0e543
Show file tree
Hide file tree
Showing 3 changed files with 89 additions and 46 deletions.
29 changes: 6 additions & 23 deletions robot-server/robot_server/persistence/_database.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@

import sqlalchemy

from server_utils import sql_utils

from ._tables import add_tables_to_db
from ._migrations import migrate

Expand All @@ -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://<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
76 changes: 76 additions & 0 deletions server-utils/server_utils/sql_utils.py
Original file line number Diff line number Diff line change
@@ -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")
30 changes: 7 additions & 23 deletions system-server/system_server/persistence/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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

0 comments on commit 5f0e543

Please sign in to comment.