Skip to content

fix: unbreak unit suite — stale persist stubs and expression-index drift flapping #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Eliebibs wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix: unbreak unit suite — stale persist stubs and expression-index drift flapping #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions app/db/migrate.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@
"idx_logs_status_error_time",
"idx_logs_api_key_time_account",
"idx_logs_source_requested_at",
"idx_logs_session_time",
}
),
"account_limit_warmups": frozenset(
Expand Down Expand Up @@ -570,11 +571,26 @@ def _unwrap_schema_drift_diff(diff: object) -> object:
return diff


def _is_manually_verified_index_diff(diff: tuple) -> bool:
# Expression-based indexes get approximate autogen signatures on some
# alembic/dialect combinations, flapping as remove_index/add_index pairs.
# Their presence is asserted by _manual_schema_drift_diffs instead.
if diff[0] not in {"add_index", "remove_index"} or len(diff) < 2:
return False
index = diff[1]
table_name = getattr(getattr(index, "table", None), "name", None)
required = _MANUAL_DRIFT_INDEX_REQUIREMENTS.get(str(table_name), frozenset())
return str(getattr(index, "name", None)) in required


def _is_ignored_schema_drift(connection: Connection, diff: object) -> bool:
diff = _unwrap_schema_drift_diff(diff)
if not isinstance(diff, tuple) or not diff:
return False

if _is_manually_verified_index_diff(diff):
return True

if diff[0] == "remove_column" and len(diff) >= 4:
column = diff[3]
column_name = getattr(column, "name", None)
Expand Down
48 changes: 37 additions & 11 deletions tests/unit/test_proxy_load_balancer_refresh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1421,14 +1421,18 @@ async def slow_persist_selection_state(
accounts_repo: AccountsRepository,
account_map: dict[str, Account],
states: list[Any],
) -> None:
*,
skip_account_ids: frozenset[str] = frozenset(),
) -> set[str]:
nonlocal inflight_persist_calls
inflight_persist_calls += 1
try:
if inflight_persist_calls >= 2:
overlap_observed.set()
await asyncio.sleep(0.05)
await original_persist_selection_state(self, accounts_repo, account_map, states)
return await original_persist_selection_state(
self, accounts_repo, account_map, states, skip_account_ids=skip_account_ids
)
finally:
inflight_persist_calls -= 1

Expand Down Expand Up @@ -1490,13 +1494,17 @@ async def controlled_persist_selection_state(
accounts_repo: AccountsRepository,
account_map: dict[str, Account],
states: list[Any],
) -> None:
*,
skip_account_ids: frozenset[str] = frozenset(),
) -> set[str]:
nonlocal blocked_once
if not blocked_once and any(state.error_count == 0 for state in states):
blocked_once = True
select_sync_blocked.set()
await release_select_sync.wait()
await original_persist_selection_state(self, accounts_repo, account_map, states)
return await original_persist_selection_state(
self, accounts_repo, account_map, states, skip_account_ids=skip_account_ids
)

monkeypatch.setattr(LoadBalancer, "_persist_selection_state", controlled_persist_selection_state)

Expand Down Expand Up @@ -1865,9 +1873,13 @@ async def test_select_account_retries_after_post_persist_permanent_failure(monke
original_persist_selection_state = balancer._persist_selection_state
injected = False

async def wrapped_persist_selection_state(accounts_repo_arg, account_map, states):
async def wrapped_persist_selection_state(
accounts_repo_arg, account_map, states, *, skip_account_ids: frozenset[str] = frozenset()
):
nonlocal injected
result = await original_persist_selection_state(accounts_repo_arg, account_map, states)
result = await original_persist_selection_state(
accounts_repo_arg, account_map, states, skip_account_ids=skip_account_ids
)
if not injected:
injected = True
await balancer.mark_permanent_failure(account, "refresh_token_expired")
Expand Down Expand Up @@ -1913,9 +1925,13 @@ async def test_select_account_retries_after_post_persist_quota_exceeded(monkeypa
original_persist_selection_state = balancer._persist_selection_state
injected = False

async def wrapped_persist_selection_state(accounts_repo_arg, account_map, states):
async def wrapped_persist_selection_state(
accounts_repo_arg, account_map, states, *, skip_account_ids: frozenset[str] = frozenset()
):
nonlocal injected
result = await original_persist_selection_state(accounts_repo_arg, account_map, states)
result = await original_persist_selection_state(
accounts_repo_arg, account_map, states, skip_account_ids=skip_account_ids
)
if not injected:
injected = True
await balancer.mark_quota_exceeded(account, {"message": "quota exceeded"})
Expand Down Expand Up @@ -2161,14 +2177,18 @@ async def stale_selected_persist(
accounts_repo: AccountsRepository,
account_map: dict[str, Account],
states: list[Any],
*,
skip_account_ids: frozenset[str] = frozenset(),
) -> set[str]:
nonlocal first_persist
if first_persist:
first_persist = False
account.status = AccountStatus.DEACTIVATED
account.deactivation_reason = "Refresh token expired - re-login required"
return {account.id}
return await original_persist_selection_state(accounts_repo, account_map, states)
return await original_persist_selection_state(
accounts_repo, account_map, states, skip_account_ids=skip_account_ids
)

monkeypatch.setattr(balancer, "_load_selection_inputs", counted_load_selection_inputs)
monkeypatch.setattr(sticky_repo, "get_account_id", pinned_account_id)
Expand Down Expand Up @@ -2244,8 +2264,10 @@ async def always_stale_selected_persist(
accounts_repo: AccountsRepository,
account_map: dict[str, Account],
states: list[Any],
*,
skip_account_ids: frozenset[str] = frozenset(),
) -> set[str]:
del accounts_repo, account_map, states
del accounts_repo, account_map, states, skip_account_ids
return {account.id}

monkeypatch.setattr(balancer, "_load_selection_inputs", counted_load_selection_inputs)
Expand Down Expand Up @@ -2663,10 +2685,14 @@ async def blocking_persist_selection_state(
accounts_repo_arg: AccountsRepository,
account_map: dict[str, Account],
states: list[Any],
*,
skip_account_ids: frozenset[str] = frozenset(),
) -> set[str]:
persist_started.set()
await release_persist.wait()
return await original_persist_selection_state(accounts_repo_arg, account_map, states)
return await original_persist_selection_state(
accounts_repo_arg, account_map, states, skip_account_ids=skip_account_ids
)

monkeypatch.setattr(balancer, "_persist_selection_state", blocking_persist_selection_state)

Expand Down