Skip to content

Commit

Permalink
Avoid mutating database schema definitions during schema migration (h…
Browse files Browse the repository at this point in the history 
…ome-assistant#122012)

* Avoid mutating database schema definitions during schema migration

* Adjust test when using mysql

* Address review comment
  • Loading branch information
@emontnemery
emontnemery authored Jul 16, 2024
1 parent c860b6c commit d8440e8
Show file tree
Hide file tree
Showing 2 changed files with 137 additions and 3 deletions.
17 changes: 15 additions & 2 deletions homeassistant/components/recorder/migration.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -576,7 +576,13 @@ def _update_states_table_with_foreign_key_options(
connection.execute(DropConstraint(alter["old_fk"])) # type: ignore[no-untyped-call]
for fkc in states_key_constraints:
if fkc.column_keys == alter["columns"]:
connection.execute(AddConstraint(fkc)) # type: ignore[no-untyped-call]
# AddConstraint mutates the constraint passed to it, we need to
# undo that to avoid changing the behavior of the table schema.
# https://github.com/sqlalchemy/sqlalchemy/blob/96f1172812f858fead45cdc7874abac76f45b339/lib/sqlalchemy/sql/ddl.py#L746-L748
create_rule = fkc._create_rule # noqa: SLF001
add_constraint = AddConstraint(fkc) # type: ignore[no-untyped-call]
fkc._create_rule = create_rule # noqa: SLF001
connection.execute(add_constraint)
except (InternalError, OperationalError):
_LOGGER.exception(
"Could not update foreign options in %s table", TABLE_STATES
Expand Down Expand Up @@ -634,10 +640,17 @@ def _restore_foreign_key_constraints(
)
continue

# AddConstraint mutates the constraint passed to it, we need to
# undo that to avoid changing the behavior of the table schema.
# https://github.com/sqlalchemy/sqlalchemy/blob/96f1172812f858fead45cdc7874abac76f45b339/lib/sqlalchemy/sql/ddl.py#L746-L748
create_rule = constraint._create_rule # noqa: SLF001
add_constraint = AddConstraint(constraint) # type: ignore[no-untyped-call]
constraint._create_rule = create_rule # noqa: SLF001

with session_scope(session=session_maker()) as session:
try:
connection = session.connection()
connection.execute(AddConstraint(constraint)) # type: ignore[no-untyped-call]
connection.execute(add_constraint)
except (InternalError, OperationalError):
_LOGGER.exception("Could not update foreign options in %s table", table)

Expand Down
123 changes: 122 additions & 1 deletion tests/components/recorder/test_migrate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import importlib
import sqlite3
import sys
from unittest.mock import Mock, PropertyMock, call, patch
from unittest.mock import ANY, Mock, PropertyMock, call, patch

import pytest
from sqlalchemy import create_engine, text
Expand Down Expand Up @@ -688,6 +688,127 @@ def test_rebuild_sqlite_states_table_extra_columns(
engine.dispose()


@pytest.mark.skip_on_db_engine(["sqlite"])
@pytest.mark.usefixtures("skip_by_db_engine")
def test_drop_restore_foreign_key_constraints(recorder_db_url: str) -> None:
"""Test we can drop and then restore foreign keys.
This is not supported on SQLite
"""

constraints_to_recreate = (
("events", "data_id"),
("states", "event_id"), # This won't be found
("states", "old_state_id"),
)

db_engine = recorder_db_url.partition("://")[0]

expected_dropped_constraints = {
"mysql": [
(
"events",
"data_id",
{
"constrained_columns": ["data_id"],
"name": ANY,
"options": {},
"referred_columns": ["data_id"],
"referred_schema": None,
"referred_table": "event_data",
},
),
(
"states",
"old_state_id",
{
"constrained_columns": ["old_state_id"],
"name": ANY,
"options": {},
"referred_columns": ["state_id"],
"referred_schema": None,
"referred_table": "states",
},
),
],
"postgresql": [
(
"events",
"data_id",
{
"comment": None,
"constrained_columns": ["data_id"],
"name": "events_data_id_fkey",
"options": {},
"referred_columns": ["data_id"],
"referred_schema": None,
"referred_table": "event_data",
},
),
(
"states",
"old_state_id",
{
"comment": None,
"constrained_columns": ["old_state_id"],
"name": "states_old_state_id_fkey",
"options": {},
"referred_columns": ["state_id"],
"referred_schema": None,
"referred_table": "states",
},
),
],
}

engine = create_engine(recorder_db_url)
db_schema.Base.metadata.create_all(engine)

with Session(engine) as session:
session_maker = Mock(return_value=session)
dropped_constraints_1 = [
dropped_constraint
for table, column in constraints_to_recreate
for dropped_constraint in migration._drop_foreign_key_constraints(
session_maker, engine, table, column
)
]
assert dropped_constraints_1 == expected_dropped_constraints[db_engine]

# Check we don't find the constrained columns again (they are removed)
with Session(engine) as session:
session_maker = Mock(return_value=session)
dropped_constraints_2 = [
dropped_constraint
for table, column in constraints_to_recreate
for dropped_constraint in migration._drop_foreign_key_constraints(
session_maker, engine, table, column
)
]
assert dropped_constraints_2 == []

# Restore the constraints
with Session(engine) as session:
session_maker = Mock(return_value=session)
migration._restore_foreign_key_constraints(
session_maker, engine, dropped_constraints_1
)

# Check we do find the constrained columns again (they are restored)
with Session(engine) as session:
session_maker = Mock(return_value=session)
dropped_constraints_3 = [
dropped_constraint
for table, column in constraints_to_recreate
for dropped_constraint in migration._drop_foreign_key_constraints(
session_maker, engine, table, column
)
]
assert dropped_constraints_3 == expected_dropped_constraints[db_engine]

engine.dispose()


def test_restore_foreign_key_constraints_with_error(
caplog: pytest.LogCaptureFixture,
) -> None:
Expand Down

0 comments on commit d8440e8

Please sign in to comment.