peewee_async_model.py

from concurrent.futures import ThreadPoolExecutor
from itertools import chain
from peewee import Model

executor = ThreadPoolExecutor(max_workers=5)


def _query_on_executor_meta(name, bases, attrs, executor):
    """A metaclass that adds async database methods to a `peewee.Model`.

    The metaclass leaves the normal database methods untouched and adds methods
    of the form '<normal method name>_async`. All async methods run on an
    executor and return a `Future`.
    """
    assert Model in bases
    class_methods = {'create', 'create_or_get', 'create_table', 'delete',
                     'drop_table', 'filter', 'get', 'get_or_create',
                     'insert', 'insert_from', 'insert_many', 'raw', 'select',
                     'sqlall', 'table_exists', 'update'}
    instance_methods = {'save', 'delete_instance', 'dependencies'}

    for method_name in chain(class_methods, instance_methods):
        def async_method(self_or_cls, *args, method_name=method_name, **kwargs):
            orig_method = getattr(self_or_cls, method_name)
            return executor.submit(orig_method, *args, **kwargs)

        if method_name in class_methods:
            async_method = classmethod(async_method)
        attrs[method_name + '_async'] = async_method
    return type(name, bases, attrs)


class AsyncModel(Model, metaclass=_query_on_executor_meta, executor=executor):
    pass