core/base/orm.py

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

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

from settings import DB_URL

if DB_URL.startswith("sqlite"):
    engine = create_engine(DB_URL)
else:
    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)


class Base(declarative_base()):
    __table__: Table
    __tablename__: str
    __new__: Callable
    __init__: Callable

    __abstract__: bool = True
    __table_args__ = {"extend_existing": True}
    id: int = 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}