core/services/db.py
Untone 435d1e4505
Some checks failed
deploy / deploy (push) Failing after 1m54s
new-version-0-2-13
2023-11-03 13:10:22 +03:00

84 lines
2.3 KiB
Python

from contextlib import contextmanager
import logging
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 sessionmaker
from sqlalchemy.sql.schema import Table
from settings import DB_URL
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
engine = create_engine(DB_URL, echo=False, pool_size=10, max_overflow=20)
Session = sessionmaker(bind=engine, expire_on_commit=False)
T = TypeVar("T")
REGISTRY: Dict[str, type] = {}
@contextmanager
def local_session():
session = Session()
try:
yield session
session.commit()
except Exception as e:
print(f"[services.db] Error session: {e}")
session.rollback()
raise
finally:
session.close()
class Base(declarative_base()):
__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]:
try:
instance = cls(**kwargs)
return instance.save()
except Exception as e:
print(f"[services.db] Error create: {e}")
return None
def save(self) -> Generic[T]:
with local_session() as session:
try:
session.add(self)
except Exception as e:
print(f"[services.db] Error save: {e}")
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)
with local_session() as session:
try:
session.commit()
except Exception as e:
print(f"[services.db] Error update: {e}")
def dict(self) -> Dict[str, Any]:
column_names = self.__table__.columns.keys()
try:
return {c: getattr(self, c) for c in column_names}
except Exception as e:
print(f"[services.db] Error dict: {e}")
return {}