Teach recorder data migrator base class to update MigrationChanges (#125214)

* Teach recorder data migrator base class to update MigrationChanges

* Bump migration version

* Improve test coverage

* Update migration.py

* Revert migrator version bump

* Remove unneeded change
This commit is contained in:
Erik Montnemery
2024-09-05 08:56:18 +02:00
committed by GitHub
parent 4c56cbe8c8
commit a8f2204f4f
2 changed files with 47 additions and 83 deletions

View File

@@ -645,23 +645,24 @@ def _is_retryable_error(instance: Recorder, err: OperationalError) -> bool:
type _FuncType[_T, **_P, _R] = Callable[Concatenate[_T, _P], _R]
type _FuncOrMethType[**_P, _R] = Callable[_P, _R]
def retryable_database_job[_RecorderT: Recorder, **_P](
description: str,
) -> Callable[[_FuncType[_RecorderT, _P, bool]], _FuncType[_RecorderT, _P, bool]]:
def retryable_database_job[**_P](
description: str, method: bool = False
) -> Callable[[_FuncOrMethType[_P, bool]], _FuncOrMethType[_P, bool]]:
"""Try to execute a database job.
The job should return True if it finished, and False if it needs to be rescheduled.
"""
recorder_pos = 1 if method else 0
def decorator(
job: _FuncType[_RecorderT, _P, bool],
) -> _FuncType[_RecorderT, _P, bool]:
def decorator(job: _FuncOrMethType[_P, bool]) -> _FuncOrMethType[_P, bool]:
@functools.wraps(job)
def wrapper(instance: _RecorderT, *args: _P.args, **kwargs: _P.kwargs) -> bool:
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> bool:
instance: Recorder = args[recorder_pos] # type: ignore[assignment]
try:
return job(instance, *args, **kwargs)
return job(*args, **kwargs)
except OperationalError as err:
if _is_retryable_error(instance, err):
assert isinstance(err.orig, BaseException) # noqa: PT017