diff --git a/homeassistant/components/recorder/migration.py b/homeassistant/components/recorder/migration.py index 324bdd5ea13480..4d9978c641b170 100644 --- a/homeassistant/components/recorder/migration.py +++ b/homeassistant/components/recorder/migration.py @@ -106,6 +106,7 @@ execute_stmt_lambda_element, get_index_by_name, retryable_database_job, + retryable_database_job_method, session_scope, ) @@ -2229,7 +2230,7 @@ def do_migrate(self, instance: Recorder, session: Session) -> None: else: self.migration_done(instance, session) - @retryable_database_job("migrate data", method=True) + @retryable_database_job_method("migrate data") def migrate_data(self, instance: Recorder) -> bool: """Migrate some data, returns True if migration is completed.""" status = self.migrate_data_impl(instance) diff --git a/homeassistant/components/recorder/util.py b/homeassistant/components/recorder/util.py index 75e403d8204283..d078c32cb88cab 100644 --- a/homeassistant/components/recorder/util.py +++ b/homeassistant/components/recorder/util.py @@ -645,44 +645,66 @@ def _is_retryable_error(instance: Recorder, err: OperationalError) -> bool: type _FuncType[**P, R] = Callable[Concatenate[Recorder, P], R] +type _MethType[Self, **P, R] = Callable[Concatenate[Self, Recorder, P], R] type _FuncOrMethType[**_P, _R] = Callable[_P, _R] def retryable_database_job[**_P]( - description: str, method: bool = False -) -> Callable[[_FuncOrMethType[_P, bool]], _FuncOrMethType[_P, bool]]: + description: str, +) -> Callable[[_FuncType[_P, bool]], _FuncType[_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: _FuncOrMethType[_P, bool]) -> _FuncOrMethType[_P, bool]: - @functools.wraps(job) - def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> bool: - instance: Recorder = args[recorder_pos] # type: ignore[assignment] - try: - return job(*args, **kwargs) - except OperationalError as err: - if _is_retryable_error(instance, err): - assert isinstance(err.orig, BaseException) # noqa: PT017 - _LOGGER.info( - "%s; %s not completed, retrying", err.orig.args[1], description - ) - time.sleep(instance.db_retry_wait) - # Failed with retryable error - return False + def decorator(job: _FuncType[_P, bool]) -> _FuncType[_P, bool]: + return _wrap_func_or_meth(job, description, False) - _LOGGER.warning("Error executing %s: %s", description, err) + return decorator - # Failed with permanent error - return True - return wrapper +def retryable_database_job_method[_Self, **_P]( + description: str, +) -> Callable[[_MethType[_Self, _P, bool]], _MethType[_Self, _P, bool]]: + """Try to execute a database job. + + The job should return True if it finished, and False if it needs to be rescheduled. + """ + + def decorator(job: _MethType[_Self, _P, bool]) -> _MethType[_Self, _P, bool]: + return _wrap_func_or_meth(job, description, True) return decorator +def _wrap_func_or_meth[**_P]( + job: _FuncOrMethType[_P, bool], description: str, method: bool +) -> _FuncOrMethType[_P, bool]: + recorder_pos = 1 if method else 0 + + @functools.wraps(job) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> bool: + instance: Recorder = args[recorder_pos] # type: ignore[assignment] + try: + return job(*args, **kwargs) + except OperationalError as err: + if _is_retryable_error(instance, err): + assert isinstance(err.orig, BaseException) # noqa: PT017 + _LOGGER.info( + "%s; %s not completed, retrying", err.orig.args[1], description + ) + time.sleep(instance.db_retry_wait) + # Failed with retryable error + return False + + _LOGGER.warning("Error executing %s: %s", description, err) + + # Failed with permanent error + return True + + return wrapper + + def database_job_retry_wrapper[**_P]( description: str, attempts: int = 5 ) -> Callable[[_FuncType[_P, None]], _FuncType[_P, None]]: