database.py

from typing import Annotated
import fakeredis
from fastapi import Depends
import redis.asyncio as redis
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from db_models import base_model
from db_models import *
import os


redis_client: redis.Redis | None = None

if os.getenv("ENVIRONMENT") == "testing":
    redis_client = fakeredis.aioredis.FakeRedis(decode_responses=True)


if os.getenv("ENVIRONMENT") == "testing":
    SQLALCHEMY_DATABASE_URL = os.getenv("TEST_DATABASE_URL")
    REDIS_URL = os.getenv("TEST_REDIS_URL")
else:
    SQLALCHEMY_DATABASE_URL = os.getenv("DATABASE_URL")
    REDIS_URL = os.getenv("REDIS_URL")


if os.getenv("ENVIRONMENT") == "production" or os.getenv("ENVIRONMENT") == "stage":
    engine = create_engine(
        SQLALCHEMY_DATABASE_URL,
        pool_size=10,
        max_overflow=20,
        pool_pre_ping=True,
        pool_recycle=1800,
    )

else:
    engine = create_engine(SQLALCHEMY_DATABASE_URL)

session_factory = sessionmaker(engine, expire_on_commit=False)


def get_db():
    with session_factory() as session:
        yield session


DB_dependency = Annotated[Session, Depends(get_db)]


async def get_redis():
    if os.getenv("ENVIRONMENT") == "testing":
        if redis_client is None:
            yield fakeredis.aioredis.FakeRedis(decode_responses=True)
        else:
            yield redis_client
    else:
        if redis_client is None:
            raise RuntimeError("Redis client not initialized")
        else:
            yield redis_client


def init_db():
    base_model.BaseModel_DB.metadata.create_all(bind=engine)