model.py 1.19 KB
import contextlib
import re

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.orm import relationship as relationship_

SQLALCHEMY_DATABASE_URL = open('./docs/postgres').read().strip()

engine = create_engine(
    SQLALCHEMY_DATABASE_URL,
    pool_size=10,
    max_overflow=5,
)
SessionLocal = sessionmaker(
    autocommit=False,
    autoflush=False,
    bind=engine,
)

Base = declarative_base()


def get_db():
    try:
        db = SessionLocal()
        yield db
    finally:
        db.close()


@contextlib.contextmanager
def transaction(session):
    if not session.in_transaction():
        with session.begin():
            yield
    else:
        yield

    session.commit()


def relationship(*args, **kwargs):
    if 'passive_deletes' not in kwargs and 'viewonly' not in kwargs:
        kwargs['passive_deletes'] = True
    return relationship_(*args, **kwargs)


def update_from_schema(model, schema):
    for field, field_data in schema.dict().items():
        setattr(model, field, field_data)
    return model


def model_columns(model):
    return [col.label(col.key) for col in model.__table__.c]