34 lines
994 B
Python
34 lines
994 B
Python
"""Database connection and session management."""
|
|
from sqlmodel import SQLModel, create_engine, Session
|
|
from app import models # Ensure all models are loaded for metadata
|
|
from typing import Generator
|
|
|
|
from app.core.config import settings
|
|
|
|
|
|
# Create database engine
|
|
# Convert asyncpg URL to psycopg2 for synchronous operations
|
|
database_url = settings.get_database_url
|
|
if "+asyncpg" in database_url:
|
|
database_url = database_url.replace("+asyncpg", "+psycopg2")
|
|
elif "postgresql://" in database_url and "+psycopg2" not in database_url:
|
|
database_url = database_url.replace("postgresql://", "postgresql+psycopg2://")
|
|
|
|
engine = create_engine(
|
|
database_url,
|
|
echo=settings.debug,
|
|
future=True,
|
|
)
|
|
|
|
|
|
def init_db() -> None:
|
|
"""Initialize database by creating all tables."""
|
|
SQLModel.metadata.create_all(engine)
|
|
|
|
|
|
def get_session() -> Generator[Session, None, None]:
|
|
"""Dependency for getting database session."""
|
|
with Session(engine) as session:
|
|
yield session
|
|
|