Skip to content

Commit

Permalink
Ensure legacy event foreign key is removed from the states table when…
Browse files Browse the repository at this point in the history
… a previous rebuild failed (#123388)

* Ensure legacy event foreign key is removed from the states table

If the system ran out of disk space removing the FK, it would
fail. #121938 fixed that to try again, however that PR was made
ineffective by #122069 since it will never reach the check.

To solve this, the migration version is incremented to 2, and
the migration is no longer marked as done unless the rebuild
/fk removal is successful.

* fix logic for mysql

* fix test

* asserts

* coverage

* coverage

* narrow test

* fixes

* split tests

* should have skipped

* fixture must be used
  • Loading branch information
bdraco authored Aug 9, 2024
1 parent 03ba8f6 commit 00c1a3f
Show file tree
Hide file tree
Showing 3 changed files with 368 additions and 16 deletions.
24 changes: 15 additions & 9 deletions homeassistant/components/recorder/migration.py
Original file line number Diff line number Diff line change
Expand Up @@ -632,7 +632,7 @@ def _update_states_table_with_foreign_key_options(

def _drop_foreign_key_constraints(
session_maker: Callable[[], Session], engine: Engine, table: str, column: str
) -> list[tuple[str, str, ReflectedForeignKeyConstraint]]:
) -> tuple[bool, list[tuple[str, str, ReflectedForeignKeyConstraint]]]:
"""Drop foreign key constraints for a table on specific columns."""
inspector = sqlalchemy.inspect(engine)
dropped_constraints = [
Expand All @@ -649,6 +649,7 @@ def _drop_foreign_key_constraints(
if foreign_key["name"] and foreign_key["constrained_columns"] == [column]
]

fk_remove_ok = True
for drop in drops:
with session_scope(session=session_maker()) as session:
try:
Expand All @@ -660,8 +661,9 @@ def _drop_foreign_key_constraints(
TABLE_STATES,
column,
)
fk_remove_ok = False

return dropped_constraints
return fk_remove_ok, dropped_constraints


def _restore_foreign_key_constraints(
Expand Down Expand Up @@ -1481,7 +1483,7 @@ def _apply_update(self) -> None:
for column in columns
for dropped_constraint in _drop_foreign_key_constraints(
self.session_maker, self.engine, table, column
)
)[1]
]
_LOGGER.debug("Dropped foreign key constraints: %s", dropped_constraints)

Expand Down Expand Up @@ -1956,14 +1958,15 @@ def cleanup_legacy_states_event_ids(instance: Recorder) -> bool:
if instance.dialect_name == SupportedDialect.SQLITE:
# SQLite does not support dropping foreign key constraints
# so we have to rebuild the table
rebuild_sqlite_table(session_maker, instance.engine, States)
fk_remove_ok = rebuild_sqlite_table(session_maker, instance.engine, States)
else:
_drop_foreign_key_constraints(
fk_remove_ok, _ = _drop_foreign_key_constraints(
session_maker, instance.engine, TABLE_STATES, "event_id"
)
_drop_index(session_maker, "states", LEGACY_STATES_EVENT_ID_INDEX)
instance.use_legacy_events_index = False
_mark_migration_done(session, EventIDPostMigration)
if fk_remove_ok:
_drop_index(session_maker, "states", LEGACY_STATES_EVENT_ID_INDEX)
instance.use_legacy_events_index = False
_mark_migration_done(session, EventIDPostMigration)

return True

Expand Down Expand Up @@ -2419,6 +2422,7 @@ class EventIDPostMigration(BaseRunTimeMigration):

migration_id = "event_id_post_migration"
task = MigrationTask
migration_version = 2

@staticmethod
def migrate_data(instance: Recorder) -> bool:
Expand Down Expand Up @@ -2469,7 +2473,7 @@ def _mark_migration_done(

def rebuild_sqlite_table(
session_maker: Callable[[], Session], engine: Engine, table: type[Base]
) -> None:
) -> bool:
"""Rebuild an SQLite table.
This must only be called after all migrations are complete
Expand Down Expand Up @@ -2524,8 +2528,10 @@ def rebuild_sqlite_table(
# Swallow the exception since we do not want to ever raise
# an integrity error as it would cause the database
# to be discarded and recreated from scratch
return False
else:
_LOGGER.warning("Rebuilding SQLite table %s finished", orig_name)
return True
finally:
with session_scope(session=session_maker()) as session:
# Step 12 - Re-enable foreign keys
Expand Down
14 changes: 7 additions & 7 deletions tests/components/recorder/test_migrate.py
Original file line number Diff line number Diff line change
Expand Up @@ -748,7 +748,7 @@ def test_rebuild_sqlite_states_table(recorder_db_url: str) -> None:
session.add(States(state="on"))
session.commit()

migration.rebuild_sqlite_table(session_maker, engine, States)
assert migration.rebuild_sqlite_table(session_maker, engine, States) is True

with session_scope(session=session_maker()) as session:
assert session.query(States).count() == 1
Expand Down Expand Up @@ -776,13 +776,13 @@ def test_rebuild_sqlite_states_table_missing_fails(
session.connection().execute(text("DROP TABLE states"))
session.commit()

migration.rebuild_sqlite_table(session_maker, engine, States)
assert migration.rebuild_sqlite_table(session_maker, engine, States) is False
assert "Error recreating SQLite table states" in caplog.text
caplog.clear()

# Now rebuild the events table to make sure the database did not
# get corrupted
migration.rebuild_sqlite_table(session_maker, engine, Events)
assert migration.rebuild_sqlite_table(session_maker, engine, Events) is True

with session_scope(session=session_maker()) as session:
assert session.query(Events).count() == 1
Expand Down Expand Up @@ -812,7 +812,7 @@ def test_rebuild_sqlite_states_table_extra_columns(
text("ALTER TABLE states ADD COLUMN extra_column TEXT")
)

migration.rebuild_sqlite_table(session_maker, engine, States)
assert migration.rebuild_sqlite_table(session_maker, engine, States) is True
assert "Error recreating SQLite table states" not in caplog.text

with session_scope(session=session_maker()) as session:
Expand Down Expand Up @@ -905,7 +905,7 @@ def test_drop_restore_foreign_key_constraints(recorder_db_url: str) -> None:
for table, column in constraints_to_recreate
for dropped_constraint in migration._drop_foreign_key_constraints(
session_maker, engine, table, column
)
)[1]
]
assert dropped_constraints_1 == expected_dropped_constraints[db_engine]

Expand All @@ -917,7 +917,7 @@ def test_drop_restore_foreign_key_constraints(recorder_db_url: str) -> None:
for table, column in constraints_to_recreate
for dropped_constraint in migration._drop_foreign_key_constraints(
session_maker, engine, table, column
)
)[1]
]
assert dropped_constraints_2 == []

Expand All @@ -936,7 +936,7 @@ def test_drop_restore_foreign_key_constraints(recorder_db_url: str) -> None:
for table, column in constraints_to_recreate
for dropped_constraint in migration._drop_foreign_key_constraints(
session_maker, engine, table, column
)
)[1]
]
assert dropped_constraints_3 == expected_dropped_constraints[db_engine]

Expand Down
Loading

0 comments on commit 00c1a3f

Please sign in to comment.