Source code for caspia.homeserver.models.db

import asyncio
import functools
from concurrent.futures import ThreadPoolExecutor

_db_executor = ThreadPoolExecutor(max_workers=1)
_db_instance = None


def make_async():
    def decorator(func):
        @functools.wraps(func)
        async def async_version(*args, loop=None, **kwargs):
            loop = loop or asyncio.get_event_loop()
            target = functools.partial(func, *args, **kwargs)
            return await loop.run_in_executor(_db_executor, target)

        return async_version

    return decorator


def register_db(db):
    global _db_instance
    _db_instance = db


def get_db():
    if not _db_instance:
        raise RuntimeWarning('Database not set. Use `register_db`.')
    return _db_instance