test_postgres_restore_drop_cascade.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163
  1. """Regression test for the Postgres restore drop-tables-with-CASCADE fix.
  2. The bug: the restore path called `metadata.drop_all`, which only drops
  3. tables defined in the SQLAlchemy ORM and emits plain `DROP TABLE` (no
  4. CASCADE). When the live DB carries orphan tables from removed features
  5. (e.g. legacy `spoolman_slot_assignments` whose `_printer_id_fkey`
  6. constraint still references `printers`), Postgres refuses with
  7. `DependentObjectsStillExistError` and the entire restore aborts before
  8. any rows land.
  9. The fix: drop every table in the `public` schema with `CASCADE` via a
  10. `pg_tables`-iterating PL/pgSQL `DO` block, then re-create from the
  11. ORM metadata. CASCADE removes external constraints alongside the table,
  12. so orphan tables can no longer block the restore.
  13. These tests guard against a regression to `metadata.drop_all` (which
  14. would re-introduce the bug for any user with orphan tables).
  15. """
  16. from __future__ import annotations
  17. import sqlite3
  18. import tempfile
  19. from pathlib import Path
  20. from unittest.mock import AsyncMock, MagicMock, patch
  21. import pytest
  22. def _make_sqlite_source() -> Path:
  23. """Build a tiny SQLite file with one ORM-known table so the restore
  24. function progresses past its `tables_to_import & metadata.tables` gate."""
  25. with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp:
  26. path = Path(tmp.name)
  27. conn = sqlite3.connect(str(path))
  28. # `users` is in the ORM metadata so `tables_to_import` is non-empty.
  29. conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, username TEXT)")
  30. conn.commit()
  31. conn.close()
  32. return path
  33. @pytest.mark.asyncio
  34. async def test_restore_drops_tables_with_cascade_not_metadata_drop_all():
  35. """Verify the restore drop phase issues a CASCADE-aware DROP TABLE
  36. iteration over `public` schema rather than `metadata.drop_all`.
  37. Regression: prior to the fix, an orphan table holding an FK back to
  38. `printers` (e.g. legacy `spoolman_slot_assignments_printer_id_fkey`)
  39. would cause `metadata.drop_all` to fail with
  40. `DependentObjectsStillExistError`, aborting the whole restore."""
  41. from backend.app.api.routes import settings as settings_module
  42. sqlite_path = _make_sqlite_source()
  43. try:
  44. executed_sql: list[str] = []
  45. run_sync_calls: list[str] = []
  46. # Capture the exact SQL emitted on the Postgres connection.
  47. mock_conn = MagicMock()
  48. mock_conn.execute = AsyncMock(
  49. side_effect=lambda stmt, *a, **k: executed_sql.append(getattr(stmt, "text", str(stmt)))
  50. )
  51. # `await conn.run_sync(metadata.create_all)` is the only run_sync
  52. # the fix should issue. `metadata.drop_all` must never appear.
  53. async def _run_sync(fn, *args, **kw):
  54. name = getattr(fn, "__name__", repr(fn))
  55. run_sync_calls.append(name)
  56. return None
  57. mock_conn.run_sync = AsyncMock(side_effect=_run_sync)
  58. # `pg_engine.begin()` is used twice (drop+create, then import).
  59. # Both must yield the same captured-conn so we observe everything.
  60. begin_cm = MagicMock()
  61. begin_cm.__aenter__ = AsyncMock(return_value=mock_conn)
  62. begin_cm.__aexit__ = AsyncMock(return_value=False)
  63. mock_engine = MagicMock()
  64. mock_engine.begin = MagicMock(return_value=begin_cm)
  65. mock_engine.dispose = AsyncMock()
  66. # `_create_engine` is imported lazily inside the function via
  67. # `from backend.app.core.database import ... _create_engine`,
  68. # so we patch the module it's imported FROM, not settings.py.
  69. with patch(
  70. "backend.app.core.database._create_engine",
  71. new=MagicMock(return_value=mock_engine),
  72. ):
  73. await settings_module._import_sqlite_to_postgres(sqlite_path, "postgresql+asyncpg://test/test")
  74. # 1. CASCADE drop is emitted, hitting every public-schema table.
  75. cascade_drops = [s for s in executed_sql if "CASCADE" in s and "pg_tables" in s]
  76. assert cascade_drops, (
  77. "Expected a CASCADE-aware DROP TABLE iteration over the public "
  78. "schema in the restore SQL stream. Without it, orphan tables "
  79. "with FK constraints back to ORM tables (e.g. legacy "
  80. "spoolman_slot_assignments) abort the restore. Captured SQL: " + "; ".join(s[:120] for s in executed_sql)
  81. )
  82. # 2. The DO block iterates pg_tables (not just one DROP) so every
  83. # table is handled, including orphan ones not in the ORM.
  84. do_block = cascade_drops[0]
  85. assert "DROP TABLE" in do_block
  86. assert "schemaname = 'public'" in do_block
  87. # 3. `metadata.drop_all` is never invoked — that was the buggy
  88. # path. `metadata.create_all` is fine; it rebuilds the schema
  89. # after the CASCADE drop.
  90. assert "drop_all" not in run_sync_calls, (
  91. f"metadata.drop_all should not be called (regression): {run_sync_calls}"
  92. )
  93. assert "create_all" in run_sync_calls, f"metadata.create_all should still be called: {run_sync_calls}"
  94. # 4. Drop runs before create. The captured SQL is in execution order
  95. # within the same pg_engine.begin() block, and run_sync_calls is
  96. # in invocation order across both blocks.
  97. first_create_idx = run_sync_calls.index("create_all")
  98. # No drop_all anywhere — the cascade DO block (executed via .execute,
  99. # not run_sync) is what runs first. Its presence is confirmed above.
  100. assert first_create_idx >= 0
  101. finally:
  102. sqlite_path.unlink(missing_ok=True)
  103. @pytest.mark.asyncio
  104. async def test_restore_cascade_drop_targets_only_public_schema():
  105. """Defensive: the CASCADE drop must scope to the `public` schema so a
  106. shared Postgres holding non-Bambuddy tables in other schemas doesn't
  107. lose data on restore."""
  108. from backend.app.api.routes import settings as settings_module
  109. sqlite_path = _make_sqlite_source()
  110. try:
  111. executed_sql: list[str] = []
  112. mock_conn = MagicMock()
  113. mock_conn.execute = AsyncMock(
  114. side_effect=lambda stmt, *a, **k: executed_sql.append(getattr(stmt, "text", str(stmt)))
  115. )
  116. mock_conn.run_sync = AsyncMock()
  117. begin_cm = MagicMock()
  118. begin_cm.__aenter__ = AsyncMock(return_value=mock_conn)
  119. begin_cm.__aexit__ = AsyncMock(return_value=False)
  120. mock_engine = MagicMock()
  121. mock_engine.begin = MagicMock(return_value=begin_cm)
  122. mock_engine.dispose = AsyncMock()
  123. with patch(
  124. "backend.app.core.database._create_engine",
  125. new=MagicMock(return_value=mock_engine),
  126. ):
  127. await settings_module._import_sqlite_to_postgres(sqlite_path, "postgresql+asyncpg://test/test")
  128. cascade = next((s for s in executed_sql if "CASCADE" in s), None)
  129. assert cascade is not None
  130. # Schema scope check: we're not iterating `pg_class` /
  131. # `information_schema.tables` without a schema filter, which
  132. # would catch system catalogs or other-app tables.
  133. assert "schemaname = 'public'" in cascade, f"CASCADE drop must filter to public schema; got: {cascade[:200]}"
  134. assert "schemaname = '*'" not in cascade
  135. finally:
  136. sqlite_path.unlink(missing_ok=True)