database.py

from sqlalchemy import Column, DateTime, create_engine, func, text
from sqlalchemy.engine.url import URL
from sqlalchemy.ext.declarative import declared_attr

# We need to ignore type checks on this line due to https://github.com/dropbox/sqlalchemy-stubs/issues/250
from sqlalchemy.orm import as_declarative, sessionmaker  # type: ignore

from service.core.settings import settings

database_url = URL(
    "postgresql",
    host=settings.PGHOST,
    username=settings.PGUSER,
    password=settings.PGPASSWORD,
    database=settings.PGDATABASE,
    port=settings.PGPORT,
    query={},
)

engine = create_engine(database_url)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)


@as_declarative()
class Base:
    @declared_attr
    def created_at(cls):
        return Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )

    @declared_attr
    def updated_at(cls):
        return Column(
            DateTime(timezone=True), nullable=False, server_default=func.now()
        )


def get_db():
    db = SessionLocal()
    db.execute(text(f"set search_path to {settings.APP_SCHEMA}"))
    try:
        yield db
    finally:
        db.close()