royalnet/royalnet/database/alchemy.py

import typing
import asyncio
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from contextlib import contextmanager, asynccontextmanager
from ..utils import cdj

loop = asyncio.get_event_loop()


class Alchemy:
    def __init__(self, database_uri: str = "sqlite://", tables: typing.Optional[typing.Set] = None):
        self.engine = create_engine(database_uri)
        self.Base = declarative_base(bind=self.engine)
        self.Session = sessionmaker(bind=self.engine)
        self._create_tables(tables)

    def _create_tables(self, tables: typing.Optional[typing.List]):
        for table in tables:
            name = table.__name__
            try:
                self.__getattribute__(name)
            except AttributeError:
                # Actually the intended result
                self.__setattr__(name, type(name, (self.Base,), cdj(table)))
            else:
                raise NameError(f"{name} is a reserved name and can't be used as a table name")
        self.Base.metadata.create_all()

    @contextmanager
    async def session_cm(self):
        session = self.Session()
        try:
            yield session
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()

    @asynccontextmanager
    async def session_acm(self):
        session = await loop.run_in_executor(None, self.Session)
        try:
            yield session
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()