from sqlalchemy import create_engine from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker import os from pathlib import Path PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent USE_POSTGRES = os.getenv('USE_POSTGRES', 'true').lower() == 'true' if USE_POSTGRES: POSTGRES_CONFIG = { 'host': os.getenv('PG_HOST', '127.0.0.1'), 'port': os.getenv('PG_PORT', '5432'), 'database': os.getenv('PG_DATABASE', 'yzr_nr'), 'user': os.getenv('PG_USER', 'yzr_nr'), 'password': os.getenv('PG_PASSWORD', 'aTX3WKKnPfRnM5PC') } SQLALCHEMY_DATABASE_URL = ( f"postgresql://{POSTGRES_CONFIG['user']}:{POSTGRES_CONFIG['password']}" f"@{POSTGRES_CONFIG['host']}:{POSTGRES_CONFIG['port']}/{POSTGRES_CONFIG['database']}" ) engine = create_engine( SQLALCHEMY_DATABASE_URL, pool_pre_ping=True, pool_size=10, max_overflow=20 ) else: DATA_DIR = os.getenv('DATA_DIR', str(PROJECT_ROOT / 'data')) os.makedirs(DATA_DIR, exist_ok=True) DB_PATH = os.path.join(DATA_DIR, 'yzr.db') SQLALCHEMY_DATABASE_URL = f"sqlite:///{DB_PATH}" engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() def init_db(): Base.metadata.create_all(bind=engine) def get_db(): db = SessionLocal() try: yield db finally: db.close()