What is the best way to implement a context_task for periodic_tasks in Huey?

Question

I'm looking for the best way to implement Peewee's context manager with periodic tasks in Huey. Normal tasks have that nice little Huey.context_task() decorator, but there doesn't seem to be anything similar for periodic tasks.

Am I correct to assume I will just have to use an (uglier) with statement within periodic tasks?

Answer 1

Should be able to do something like this:

from functools import wraps

huey = RedisHuey()
db = PostgresqlDatabase(...)

def db_periodic_task(*args, **kwargs):
    def decorator(fn):
        @wraps(fn)
        def new_task():
            with db:
                fn()
        return huey.periodic_task(*args, **kwargs)(new_task)
    return decorator

@db_periodic_task(crontab('0', '0'))
def my_periodic_task():
    # db will have been opened at start.
    # db will then be closed upon return.