74 lines
1.8 KiB
Python
74 lines
1.8 KiB
Python
"""
|
|
Database connection and session management
|
|
"""
|
|
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
|
from sqlalchemy.orm import declarative_base
|
|
from app.config import settings
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Create async engine
|
|
engine = create_async_engine(
|
|
settings.database_url,
|
|
echo=False, # Set to True for SQL query logging
|
|
pool_size=settings.database_pool_size,
|
|
max_overflow=settings.database_max_overflow,
|
|
pool_pre_ping=True, # Verify connections before using
|
|
)
|
|
|
|
# Create async session factory
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
autocommit=False,
|
|
autoflush=False,
|
|
)
|
|
|
|
# Base class for ORM models
|
|
Base = declarative_base()
|
|
|
|
|
|
async def get_db() -> AsyncSession:
|
|
"""
|
|
Dependency function for FastAPI to get database sessions
|
|
|
|
Usage:
|
|
@app.get("/items")
|
|
async def read_items(db: AsyncSession = Depends(get_db)):
|
|
...
|
|
"""
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
yield session
|
|
await session.commit()
|
|
except Exception:
|
|
await session.rollback()
|
|
raise
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
async def init_db():
|
|
"""Initialize database - create all tables"""
|
|
from app.models.db import (
|
|
CelestialBody,
|
|
Position,
|
|
Resource,
|
|
StaticData,
|
|
NasaCache,
|
|
)
|
|
|
|
async with engine.begin() as conn:
|
|
# Create all tables
|
|
await conn.run_sync(Base.metadata.create_all)
|
|
|
|
logger.info("Database tables created successfully")
|
|
|
|
|
|
async def close_db():
|
|
"""Close database connections"""
|
|
await engine.dispose()
|
|
logger.info("Database connections closed")
|