"""Job utilities."""

import time

import config
from json_logger import logger


def retry(custom_handler=None):
    """Retry on db errors.

    Args:
        custom_handler: Custom error handling function.

    """

    def inner(f):
        def wrapped(*args, **kwargs):
            for i in range(1, config.RETRY_COUNT + 1):
                try:
                    return f(*args, **kwargs)
                except Exception as ex:
                    if i != config.RETRY_COUNT:
                        logger.info(f"Retry on error: {ex}")
                        if custom_handler:
                            custom_handler(*args, **kwargs)
                        time.sleep(i * config.RETRY_WAIT_RATE)
                    else:
                        raise

        return wrapped

    return inner