core/base/orm.py

from typing import Any, Callable, Dict, Generic, TypeVar

from sqlalchemy import Column, Integer, create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import Session
from sqlalchemy.sql.schema import Table

from settings import DB_URL

engine = create_engine(DB_URL, echo=False, pool_size=10, max_overflow=20)

T = TypeVar("T")

REGISTRY: Dict[str, type] = {}


def local_session():
    return Session(bind=engine, expire_on_commit=False)


DeclarativeBase = declarative_base()  # type: Any


class Base(DeclarativeBase):
    __table__: Table
    __tablename__: str
    __new__: Callable
    __init__: Callable
    __allow_unmapped__ = True
    __abstract__ = True
    __table_args__ = {"extend_existing": True}

    id = Column(Integer, primary_key=True)

    def __init_subclass__(cls, **kwargs):
        REGISTRY[cls.__name__] = cls

    @classmethod
    def create(cls: Generic[T], **kwargs) -> Generic[T]:
        instance = cls(**kwargs)
        return instance.save()

    def save(self) -> Generic[T]:
        with local_session() as session:
            session.add(self)
            session.commit()
        return self

    def update(self, input):
        column_names = self.__table__.columns.keys()
        for name, value in input.items():
            if name in column_names:
                setattr(self, name, value)

    def dict(self) -> Dict[str, Any]:
        column_names = self.__table__.columns.keys()
        return {c: getattr(self, c) for c in column_names}