ENG-2882 Update post_upgrade_index_creation script to be more reliable#7583
Open
ENG-2882 Update post_upgrade_index_creation script to be more reliable#7583
Conversation
…y checks in index creation script
Contributor
|
The latest updates on your projects. Learn more about Vercel for GitHub.
1 Skipped Deployment
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughRenames Changes
Sequence DiagramssequenceDiagram
participant Backfill as Backfill Process
participant Registry as Migration Task Registry
participant DB as Database
participant Worker as Worker/Processor
Backfill->>Registry: register_backfill_started(name)
Registry->>DB: INSERT (key, task_type, completed_at=NULL)
DB-->>Registry: Inserted
Registry-->>Backfill: Confirmed
Backfill->>DB: SELECT completed_at WHERE key
DB-->>Backfill: completed_at NULL (in-progress) / NOT NULL (done)
Backfill->>Worker: Execute backfill batches
Worker->>Worker: Process data
Worker-->>Backfill: Batch complete
Backfill->>Registry: mark_backfill_completed(name)
Registry->>DB: UPDATE completed_at = NOW()
DB-->>Registry: Updated
Registry-->>Backfill: Marked complete
sequenceDiagram
participant IndexTask as Index Creation Task
participant Registry as Migration Task Registry
participant DB as Database
participant PG as PostgreSQL
IndexTask->>Registry: get_registered_index_keys()
Registry->>DB: SELECT key WHERE task_type='index' AND completed_at IS NULL
DB-->>Registry: Registered keys
Registry-->>IndexTask: Return keys
loop For each object in TABLE_OBJECT_MAP
IndexTask->>IndexTask: is_registered(object, keys)?
alt Not registered
IndexTask->>IndexTask: Skip object
else Registered
IndexTask->>PG: Check existing object/constraint
alt Already exists
IndexTask->>IndexTask: Skip creation
else Not existing
IndexTask->>PG: CREATE INDEX/CONSTRAINT
PG-->>IndexTask: Created
IndexTask->>Registry: mark_index_completed(migration_key)
Registry->>DB: UPDATE completed_at = NOW()
DB-->>Registry: Updated
end
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 3✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Comment |
erosselli
changed the title
Contributor
Greptile SummaryThis PR addresses a race condition between Alembic migrations and background index-creation scripts by introducing a registration-gating mechanism: the background script now only creates indexes whose Key changes:
Issues found:
Confidence Score: 3/5
Last reviewed commit: dcfa802 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
tests/ops/migration_tests/backfill_scripts.py/test_backfill_utils.py
Outdated
Show resolved
Hide resolved
Comment on lines
+147
to
+155
| def get_registered_index_keys(db: Session) -> set[str]: | ||
| """Return index task keys registered by migrations.""" | ||
| result = db.execute( | ||
| text( | ||
| "SELECT key FROM post_upgrade_background_migration_tasks " | ||
| "WHERE task_type = 'index'" | ||
| ) | ||
| ) | ||
| return {row[0] for row in result} |
Contributor
There was a problem hiding this comment.
get_registered_index_keys fetches all index rows including already-completed ones
get_registered_index_keys returns every row with task_type = 'index', regardless of whether completed_at is set. On every server start, check_and_create_objects calls this function and then calls should_skip_existing (a pg_indexes / pg_constraint catalog query) for every object whose key appears in the set — even objects that were already fully processed in a previous boot.
For deployments with many deferred index entries this adds unnecessary catalog queries on every startup. Filtering to only keys with completed_at IS NULL would limit the work to genuinely pending entries:
def get_registered_index_keys(db: Session) -> set[str]:
"""Return index task keys registered by migrations that have not yet completed."""
result = db.execute(
text(
"SELECT key FROM post_upgrade_background_migration_tasks "
"WHERE task_type = 'index' AND completed_at IS NULL"
)
)
return {row[0] for row in result}Note: the should_skip_existing guard still ensures correctness — this is purely a performance concern, not a correctness issue.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (4)
tests/api/migrations/test_post_upgrade_index_creation.py (2)
81-107: Consider adding a mock formark_index_completedto verify completion tracking.The test verifies object creation but doesn't verify that
mark_index_completed()is called after successful creation. Adding this assertion would ensure the completion tracking integration is tested.💡 Suggested test enhancement
+ `@patch`("fides.api.migrations.post_upgrade_index_creation.mark_index_completed") `@patch`("fides.api.migrations.post_upgrade_index_creation.get_registered_index_keys") `@patch`("fides.api.migrations.post_upgrade_index_creation.check_object_exists") `@patch`("fides.api.migrations.post_upgrade_index_creation.create_object") def test_processes_entry_when_migration_key_is_registered( - self, mock_create_object, mock_check_object_exists, mock_get_registered_keys + self, mock_create_object, mock_check_object_exists, mock_get_registered_keys, mock_mark_completed ): """Entry with registered migration_key is processed normally.""" # ... existing test code ... mock_create_object.assert_called_once() + mock_mark_completed.assert_called_once_with(MagicMock(), "ix_test_index") assert result == {"ix_test_index": "created"}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/test_post_upgrade_index_creation.py` around lines 81 - 107, The test should also mock and assert that mark_index_completed is invoked after a successful create; add a `@patch`("fides.api.migrations.post_upgrade_index_creation.mark_index_completed") parameter to the test signature, set it up as a MagicMock, then after calling check_and_create_objects(MagicMock(), table_map, mock_lock) assert mark_index_completed.assert_called_once_with("ix_test_index") (or the expected args) to verify completion tracking is invoked after create_object in the test_processes_entry_when_migration_key_is_registered test.
51-127: Consider adding a test case for existing objects being skipped.The test class covers unregistered keys and registered keys with non-existent objects, but doesn't test the case where
check_object_existsreturnsTrue(object already exists). This would complete the coverage ofcheck_and_create_objects()logic paths.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/test_post_upgrade_index_creation.py` around lines 51 - 127, Add a new test in TestCheckAndCreateObjectsMigrationKeyFiltering that covers the case when an object already exists: call check_and_create_objects with mock_get_registered_keys.return_value containing "ix_test_index" and mock_check_object_exists.return_value = True, then assert mock_create_object.assert_not_called() and that the returned result does not include "ix_test_index" (e.g., assert "ix_test_index" not in result or result.get("ix_test_index") is None) so the existing object path in check_and_create_objects is exercised and verified.src/fides/api/alembic/migrations/versions/xx_2026_03_04_1200_e3a9f1b2c4d5_rename_backfill_history_add_task_type.py (1)
124-130: Consider using parameterized inserts for consistency.While the current approach is safe (values come from hardcoded
EXISTING_INDEX_KEYS), using parameterized queries orop.bulk_insert()would be more consistent with SQL best practices. However, this is acceptable for a migration script with static values.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/fides/api/alembic/migrations/versions/xx_2026_03_04_1200_e3a9f1b2c4d5_rename_backfill_history_add_task_type.py` around lines 124 - 130, The current raw SQL insert built as values_clause and executed via op.execute should be converted to a parameterized/bulk insert for consistency: use alembic.operations.ops or SQLAlchemy's op.bulk_insert against the post_upgrade_background_migration_tasks table (referencing EXISTING_INDEX_KEYS and the place where values_clause and op.execute are used) to insert rows with columns ('key','task_type','completed_at') and ON CONFLICT semantics handled by the database or emulate conflict-safe behavior; replace the f-string construction and op.execute call with an op.bulk_insert or parameterized insert using the table object to avoid raw string interpolation.src/fides/api/migrations/post_upgrade_index_creation.py (1)
341-355: Consider extracting a shared helper for marking tasks completed.This function duplicates the pattern from
mark_backfill_completed()inutils.py. Consider a genericmark_task_completed(db, key, task_type)helper in utils.py to reduce duplication.♻️ Suggested refactor
In
utils.py:def mark_task_completed(db: Session, key: str, task_type: str) -> None: """Mark a task as completed (set completed_at = now()).""" db.execute( text( "UPDATE post_upgrade_background_migration_tasks " "SET completed_at = now() " "WHERE key = :key AND task_type = :task_type AND completed_at IS NULL" ), {"key": key, "task_type": task_type}, ) db.commit()Then
mark_backfill_completedandmark_index_completedcan call this shared helper.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@src/fides/api/migrations/post_upgrade_index_creation.py` around lines 341 - 355, Extract the duplicated SQL+commit logic into a new helper mark_task_completed(db: Session, key: str, task_type: str) in utils.py that executes the UPDATE against post_upgrade_background_migration_tasks with WHERE key = :key AND task_type = :task_type AND completed_at IS NULL, passing {"key": key, "task_type": task_type} and then db.commit(); then update mark_index_completed (and mark_backfill_completed) to return early on falsy key and call mark_task_completed(db, migration_key, "index") (or the appropriate task_type) instead of repeating the SQL/commit code.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In
`@src/fides/api/alembic/migrations/versions/xx_2026_03_04_1200_e3a9f1b2c4d5_rename_backfill_history_add_task_type.py`:
- Around line 124-130: The current raw SQL insert built as values_clause and
executed via op.execute should be converted to a parameterized/bulk insert for
consistency: use alembic.operations.ops or SQLAlchemy's op.bulk_insert against
the post_upgrade_background_migration_tasks table (referencing
EXISTING_INDEX_KEYS and the place where values_clause and op.execute are used)
to insert rows with columns ('key','task_type','completed_at') and ON CONFLICT
semantics handled by the database or emulate conflict-safe behavior; replace the
f-string construction and op.execute call with an op.bulk_insert or
parameterized insert using the table object to avoid raw string interpolation.
In `@src/fides/api/migrations/post_upgrade_index_creation.py`:
- Around line 341-355: Extract the duplicated SQL+commit logic into a new helper
mark_task_completed(db: Session, key: str, task_type: str) in utils.py that
executes the UPDATE against post_upgrade_background_migration_tasks with WHERE
key = :key AND task_type = :task_type AND completed_at IS NULL, passing {"key":
key, "task_type": task_type} and then db.commit(); then update
mark_index_completed (and mark_backfill_completed) to return early on falsy key
and call mark_task_completed(db, migration_key, "index") (or the appropriate
task_type) instead of repeating the SQL/commit code.
In `@tests/api/migrations/test_post_upgrade_index_creation.py`:
- Around line 81-107: The test should also mock and assert that
mark_index_completed is invoked after a successful create; add a
`@patch`("fides.api.migrations.post_upgrade_index_creation.mark_index_completed")
parameter to the test signature, set it up as a MagicMock, then after calling
check_and_create_objects(MagicMock(), table_map, mock_lock) assert
mark_index_completed.assert_called_once_with("ix_test_index") (or the expected
args) to verify completion tracking is invoked after create_object in the
test_processes_entry_when_migration_key_is_registered test.
- Around line 51-127: Add a new test in
TestCheckAndCreateObjectsMigrationKeyFiltering that covers the case when an
object already exists: call check_and_create_objects with
mock_get_registered_keys.return_value containing "ix_test_index" and
mock_check_object_exists.return_value = True, then assert
mock_create_object.assert_not_called() and that the returned result does not
include "ix_test_index" (e.g., assert "ix_test_index" not in result or
result.get("ix_test_index") is None) so the existing object path in
check_and_create_objects is exercised and verified.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: aa7dd887-3434-4e45-b17a-5aa123b0341a
📒 Files selected for processing (8)
changelog/7583-rename-backfill-history-add-migration-key-gating.yamlsrc/fides/api/alembic/migrations/versions/xx_2026_03_04_1200_e3a9f1b2c4d5_rename_backfill_history_add_task_type.pysrc/fides/api/db/database.pysrc/fides/api/migrations/backfill_scripts/utils.pysrc/fides/api/migrations/post_upgrade_backfill.pysrc/fides/api/migrations/post_upgrade_index_creation.pytests/api/migrations/test_post_upgrade_index_creation.pytests/ops/migration_tests/backfill_scripts.py/test_backfill_utils.py
Contributor
Author
|
@greptile re review |
Comment on lines
+371
to
+373
| if should_skip_existing(db, object_data): | ||
| lock.reacquire() | ||
| continue |
Contributor
There was a problem hiding this comment.
mark_index_completed is not called on the should_skip_existing path. When an entry is registered (appears in registered_keys with completed_at = NULL) but the object already exists in Postgres — for example because automigrate ran the migration that creates the index directly — the should_skip_existing branch exits via continue without calling mark_index_completed. As a result, completed_at stays NULL for that row indefinitely.
Every subsequent server startup, get_registered_index_keys will return that key again, should_skip_existing will return True again, and the catalog query will fire again — forever.
The fix is to mark the task completed when the object is found to already exist:
if should_skip_existing(db, object_data):
mark_index_completed(db, object_data.get("migration_key"))
lock.reacquire()
continue
Comment on lines
+136
to
+153
| @pytest.fixture(autouse=True) | ||
| def clean_migration_tasks(self, db: Session): | ||
| db.execute( | ||
| text( | ||
| "DELETE FROM post_upgrade_background_migration_tasks " | ||
| "WHERE key LIKE 'test-%'" | ||
| ) | ||
| ) | ||
| db.commit() | ||
| yield | ||
| db.execute( | ||
| text( | ||
| "DELETE FROM post_upgrade_background_migration_tasks " | ||
| "WHERE key LIKE 'test-%'" | ||
| ) | ||
| ) | ||
| db.commit() | ||
|
|
Contributor
There was a problem hiding this comment.
The clean_migration_tasks autouse fixture uses raw DELETE statements both as setup and teardown, which violates the project rule that the database is automatically cleared between test runs and manual deletions should be avoided.
The same pattern also appears in tests/api/migrations/backfill_scripts/test_backfill_utils.py (lines 138–153).
A straightforward fix for both is to use unique key names scoped to each test (e.g., via uuid.uuid4()) to achieve test isolation without any cleanup. Since the database is already cleared between runs, manual deletion is unnecessary:
import uuid
@pytest.fixture(autouse=True)
def unique_test_prefix() -> str:
return f"test-{uuid.uuid4()}"Then use unique_test_prefix instead of a hardcoded "test-..." string inside each test, removing the need for the clean_migration_tasks fixture entirely.
Context Used: Rule from dashboard - Do not manually delete database records in test fixtures or at the end of tests, as the database is ... (source)
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/api/migrations/backfill_scripts/test_backfill_utils.py (1)
574-577:⚠️ Potential issue | 🟡 MinorLock-refresh assertion is weaker than the test contract.
Given fixed input of 25 batches, this should assert exactly 3 refreshes (batch 10, 20, and final batch 25).
>= 2can hide a missed final refresh.🧪 Tighten the assertion
- assert mock_refresh.call_count >= 2 # At least at batch 10 and 20 + assert mock_refresh.call_count == 3 + for call_args in mock_refresh.call_args_list: + assert call_args[0] == (mock_lock,)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/backfill_scripts/test_backfill_utils.py` around lines 574 - 577, The assertion on lock refresh count is too weak; update the test in test_backfill_utils.py to assert that mock_refresh.call_count == 3 (exactly three refreshes for batches 10, 20, and final batch 25) instead of assert mock_refresh.call_count >= 2; locate the assertion that references mock_refresh.call_count in the backfill test and change it to an equality check to ensure the final refresh is verified.
🧹 Nitpick comments (4)
tests/api/migrations/backfill_scripts/test_backfill_utils.py (2)
508-510: “Before batch processing” is not actually asserted yet.The current assertion proves registration happened, but not ordering vs. batch execution. Consider asserting
mock_register.calledinside themock_executeside effect.💡 Suggested ordering check
- mock_execute.side_effect = lambda fn, **kwargs: fn() + def execute_side_effect(fn, **kwargs): + assert mock_register.called, ( + "register_backfill_started() should run before first batch execution" + ) + return fn() + mock_execute.side_effect = execute_side_effect🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/backfill_scripts/test_backfill_utils.py` around lines 508 - 510, The test currently only asserts mock_register.assert_called_once_with(mock_db, "test-backfill") but doesn't verify ordering vs batch execution; update the mock for batch execution (mock_execute) to use a side effect that asserts mock_register.called is True before proceeding (e.g., in the side effect body assert mock_register.called), keeping the existing mock_register.assert_called_once_with(mock_db, "test-backfill") assertion as well; look for the test function using mock_execute and mock_register in tests/api/migrations/backfill_scripts/test_backfill_utils.py and add the ordering check inside the mock_execute side effect to ensure register ran before batch processing.
155-237: Add a non-backfilltask_typenegative case to both completion tests.Now that
post_upgrade_background_migration_tasksstores multiple task types, these tests should also assert that a completed row for a different task type does not satisfyis_backfill_completed*().💡 Suggested test extension
`@pytest.mark.parametrize`( "record_state,expected", [ ("completed", True), ("in_progress", False), ("absent", False), + ("other_task_type_completed", False), ], - ids=["completed_record", "in_progress_record", "no_record"], + ids=["completed_record", "in_progress_record", "no_record", "other_task_type_completed"], ) ... +elif record_state == "other_task_type_completed": + db.execute( + text( + "INSERT INTO post_upgrade_background_migration_tasks (key, task_type, completed_at) " + "VALUES (:name, 'deferred_index_creation', now())" + ), + {"name": backfill_name}, + ) + db.commit()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/backfill_scripts/test_backfill_utils.py` around lines 155 - 237, Add a negative case to both test_is_backfill_completed_sync and test_is_backfill_completed_async: when a row exists with the same key and a non-'backfill' task_type but has completed_at set, is_backfill_completed(db, name) and is_backfill_completed_async(session, name) should still return False; update the parametrization or add an extra branch where you INSERT INTO post_upgrade_background_migration_tasks using task_type != 'backfill' (e.g. 'other') with now() for completed_at (use db.execute/db.commit and await async_session.execute/commit respectively) and assert the result equals False, referencing the existing functions is_backfill_completed and is_backfill_completed_async and the post_upgrade_background_migration_tasks table to locate the insertion logic.tests/api/migrations/test_post_upgrade_index_creation.py (2)
137-150: Narrow fixture cleanup scope to reduce cross-test interference.Using
WHERE key LIKE 'test-%'can delete rows from unrelated tests that share the same prefix. Prefer a class-specific prefix and parameterized cleanup.Diff suggestion
class TestRegisteredIndexKeys: """Tests for index key registration functions.""" + TEST_KEY_PREFIX = "test-post-upgrade-index-" `@pytest.fixture`(autouse=True) def clean_migration_tasks(self, db: Session): db.execute( text( "DELETE FROM post_upgrade_background_migration_tasks " - "WHERE key LIKE 'test-%'" - ) + "WHERE key LIKE :prefix" + ), + {"prefix": f"{self.TEST_KEY_PREFIX}%"}, ) db.commit() yield db.execute( text( "DELETE FROM post_upgrade_background_migration_tasks " - "WHERE key LIKE 'test-%'" - ) + "WHERE key LIKE :prefix" + ), + {"prefix": f"{self.TEST_KEY_PREFIX}%"}, ) db.commit() def test_get_registered_index_keys(self, db: Session): """Verify get_registered_index_keys() returns only index-type keys.""" - test_key = "test-index-key-for-get-registered" + test_key = f"{self.TEST_KEY_PREFIX}registered" @@ - backfill_key = "test-backfill-key-for-get-registered" + backfill_key = f"{self.TEST_KEY_PREFIX}backfill"Also applies to: 156-183
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/test_post_upgrade_index_creation.py` around lines 137 - 150, The cleanup in the fixture clean_migration_tasks uses a broad WHERE key LIKE 'test-%' that can remove other tests' rows; change the deletion to target a class- or test-specific prefix (e.g., include the test class name or a unique param) and make the fixture accept or construct that prefix (so clean_migration_tasks deletes WHERE key LIKE '<unique_prefix>-%' before and after yielding); update any other similar cleanup blocks (the one around lines 156-183) to use the same class-specific/parameterized prefix to avoid cross-test interference.
87-110: Assert completion bookkeeping in the registered-key happy path.This test verifies creation, but it doesn’t verify that the migration task is marked completed. That side effect is part of the reliability contract and should be locked in here.
Diff suggestion
`@patch`("fides.api.migrations.post_upgrade_index_creation.get_registered_index_keys") `@patch`("fides.api.migrations.post_upgrade_index_creation.check_object_exists") `@patch`("fides.api.migrations.post_upgrade_index_creation.create_object") +@patch("fides.api.migrations.post_upgrade_index_creation.mark_index_completed") def test_processes_entry_when_migration_key_is_registered( - self, mock_create_object, mock_check_object_exists, mock_get_registered_keys + self, + mock_mark_index_completed, + mock_create_object, + mock_check_object_exists, + mock_get_registered_keys, ): @@ - result = check_and_create_objects(MagicMock(), table_map, mock_lock) + db = MagicMock() + result = check_and_create_objects(db, table_map, mock_lock) @@ mock_check_object_exists.assert_called_once() mock_create_object.assert_called_once() + mock_mark_index_completed.assert_called_once_with(db, "ix_test_index") assert result == {"ix_test_index": "created"}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/test_post_upgrade_index_creation.py` around lines 87 - 110, The test currently verifies creation but not that the migration task is marked completed; after calling check_and_create_objects(MagicMock(), table_map, mock_lock) add an assertion that the completion bookkeeping was invoked for migration_key "ix_test_index" (e.g., assert mock_lock.release.assert_called_once() or whichever method your locking API uses to mark completion), referencing the same mock_lock and the migration_key "ix_test_index" used in the test to ensure the post-creation completion side effect is asserted.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Outside diff comments:
In `@tests/api/migrations/backfill_scripts/test_backfill_utils.py`:
- Around line 574-577: The assertion on lock refresh count is too weak; update
the test in test_backfill_utils.py to assert that mock_refresh.call_count == 3
(exactly three refreshes for batches 10, 20, and final batch 25) instead of
assert mock_refresh.call_count >= 2; locate the assertion that references
mock_refresh.call_count in the backfill test and change it to an equality check
to ensure the final refresh is verified.
---
Nitpick comments:
In `@tests/api/migrations/backfill_scripts/test_backfill_utils.py`:
- Around line 508-510: The test currently only asserts
mock_register.assert_called_once_with(mock_db, "test-backfill") but doesn't
verify ordering vs batch execution; update the mock for batch execution
(mock_execute) to use a side effect that asserts mock_register.called is True
before proceeding (e.g., in the side effect body assert mock_register.called),
keeping the existing mock_register.assert_called_once_with(mock_db,
"test-backfill") assertion as well; look for the test function using
mock_execute and mock_register in
tests/api/migrations/backfill_scripts/test_backfill_utils.py and add the
ordering check inside the mock_execute side effect to ensure register ran before
batch processing.
- Around line 155-237: Add a negative case to both
test_is_backfill_completed_sync and test_is_backfill_completed_async: when a row
exists with the same key and a non-'backfill' task_type but has completed_at
set, is_backfill_completed(db, name) and is_backfill_completed_async(session,
name) should still return False; update the parametrization or add an extra
branch where you INSERT INTO post_upgrade_background_migration_tasks using
task_type != 'backfill' (e.g. 'other') with now() for completed_at (use
db.execute/db.commit and await async_session.execute/commit respectively) and
assert the result equals False, referencing the existing functions
is_backfill_completed and is_backfill_completed_async and the
post_upgrade_background_migration_tasks table to locate the insertion logic.
In `@tests/api/migrations/test_post_upgrade_index_creation.py`:
- Around line 137-150: The cleanup in the fixture clean_migration_tasks uses a
broad WHERE key LIKE 'test-%' that can remove other tests' rows; change the
deletion to target a class- or test-specific prefix (e.g., include the test
class name or a unique param) and make the fixture accept or construct that
prefix (so clean_migration_tasks deletes WHERE key LIKE '<unique_prefix>-%'
before and after yielding); update any other similar cleanup blocks (the one
around lines 156-183) to use the same class-specific/parameterized prefix to
avoid cross-test interference.
- Around line 87-110: The test currently verifies creation but not that the
migration task is marked completed; after calling
check_and_create_objects(MagicMock(), table_map, mock_lock) add an assertion
that the completion bookkeeping was invoked for migration_key "ix_test_index"
(e.g., assert mock_lock.release.assert_called_once() or whichever method your
locking API uses to mark completion), referencing the same mock_lock and the
migration_key "ix_test_index" used in the test to ensure the post-creation
completion side effect is asserted.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1f9877b6-38d0-4c04-a369-ff1328875ed0
📒 Files selected for processing (5)
tests/api/migrations/backfill_scripts/__init__.pytests/api/migrations/backfill_scripts/test_backfill_stagedresource_is_leaf.pytests/api/migrations/backfill_scripts/test_backfill_stagedresourceancestor_distance.pytests/api/migrations/backfill_scripts/test_backfill_utils.pytests/api/migrations/test_post_upgrade_index_creation.py
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/api/migrations/test_post_upgrade_index_creation.py (2)
68-85: Cleanup fixture logic is duplicated across classes.The setup/teardown SQL is repeated in two autouse fixtures. Extracting a shared helper fixture would reduce drift when cleanup criteria change.
Also applies to: 150-167
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/test_post_upgrade_index_creation.py` around lines 68 - 85, The autouse fixture clean_migration_tasks duplicates identical setup/teardown SQL in multiple test classes; extract that SQL into a single shared autouse fixture (e.g., cleanup_migration_tasks) that executes the DELETE ... WHERE key LIKE 'test-%' OR key = 'ix_test_index' via db.execute/text and db.commit, then have tests rely on that one fixture instead of defining clean_migration_tasks twice (remove the duplicated fixture definitions at both locations and reference the shared fixture name where needed).
98-117: Make the registered-key happy-path test deterministic.At Line 114, this test still depends on
should_skip_existing()runtime DB state. If an index with the same name already exists, this can become flaky. Consider forcing that branch toFalsein this test.Proposed test hardening
- `@patch`("fides.api.migrations.post_upgrade_index_creation.create_object") + `@patch`( + "fides.api.migrations.post_upgrade_index_creation.should_skip_existing", + return_value=False, + ) + `@patch`("fides.api.migrations.post_upgrade_index_creation.create_object") def test_processes_entry_when_migration_key_is_registered( - self, mock_create_object, db: Session + self, mock_create_object, _mock_should_skip_existing, db: Session ):🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/api/migrations/test_post_upgrade_index_creation.py` around lines 98 - 117, The test test_processes_entry_when_migration_key_is_registered is flaky because it relies on the runtime DB state inside should_skip_existing(); make the happy-path deterministic by forcing should_skip_existing to return False for this test (e.g., patch/mock the should_skip_existing function used by check_and_create_objects so it always returns False), then run check_and_create_objects(db, self.TABLE_MAP, mock_lock) and assert as before; ensure you reference and patch the should_skip_existing symbol from the module that defines check_and_create_objects so the branch that skips creation cannot execute.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@tests/api/migrations/test_post_upgrade_index_creation.py`:
- Around line 68-85: The autouse fixture clean_migration_tasks duplicates
identical setup/teardown SQL in multiple test classes; extract that SQL into a
single shared autouse fixture (e.g., cleanup_migration_tasks) that executes the
DELETE ... WHERE key LIKE 'test-%' OR key = 'ix_test_index' via db.execute/text
and db.commit, then have tests rely on that one fixture instead of defining
clean_migration_tasks twice (remove the duplicated fixture definitions at both
locations and reference the shared fixture name where needed).
- Around line 98-117: The test
test_processes_entry_when_migration_key_is_registered is flaky because it relies
on the runtime DB state inside should_skip_existing(); make the happy-path
deterministic by forcing should_skip_existing to return False for this test
(e.g., patch/mock the should_skip_existing function used by
check_and_create_objects so it always returns False), then run
check_and_create_objects(db, self.TABLE_MAP, mock_lock) and assert as before;
ensure you reference and patch the should_skip_existing symbol from the module
that defines check_and_create_objects so the branch that skips creation cannot
execute.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5d11ed89-8843-4925-87e8-3f2deac742be
📒 Files selected for processing (1)
tests/api/migrations/test_post_upgrade_index_creation.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ticket ENG-2882
Description Of Changes
Original problem
Our migration system and background index-creation scripts can conflict, causing errors on server restart. Specifically:
The core issue is that migrations and background index-creation scripts operate without awareness of each other, leading to situations where an index can be created out-of-band and then conflict with the migration that was supposed to create it.
Proposed solution
The proposed solution involves tracking the deferred indexes that need to be created in a table, and have the
post_upgrade_index_creationscript only create the indexes that have entries on that table.This PR renames
backfill_historytopost_upgrade_background_migration_tasksand generalizes it to track both backfill and deferred index creation tasks. Previously, background index creation had no gating — all entries inTABLE_OBJECT_MAPwere processed unconditionally. Now, each entry requires amigration_keyregistered in the new table by its Alembic migration before the background task will create it.Code Changes
post_upgrade_background_migration_tasksadds new columns to it, and pre-populates it with entries for all the deferred indexes created before the migration.post_upgrade_backfill.pyscript to check the new table name, and also to register entries with a NULLcompleted_atat the beginning and update thecompleted_atvalue at the end.post_upgrade_index_creation.pyscript so allTABLE_OBJECT_MAPentries contain amigration_keyand added registration gating to only create indexes that have been added topost_upgrade_background_migration_tasksSteps to Confirm
post_upgrade_background_migration_taskstable.TABLE_OBJECT_MAPPre-Merge Checklist
CHANGELOG.mdupdatedmaindowngrade()migration is correct and worksSummary by CodeRabbit