mirror of
https://github.com/wahyd4/you-music.git
synced 2026-08-16 16:46:42 +10:00
78 lines
2.7 KiB
Python
78 lines
2.7 KiB
Python
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
|
from sqlalchemy.orm import declarative_base
|
|
from sqlalchemy.pool import NullPool
|
|
from sqlalchemy import text
|
|
from app.core.config import settings
|
|
|
|
# Configure database engine based on DB type
|
|
if settings.DB_TYPE == "postgres":
|
|
# PostgreSQL configuration - use connection pooling for better performance
|
|
# Note: For async engines, we don't explicitly set poolclass
|
|
# The async engine handles pooling internally
|
|
engine = create_async_engine(
|
|
settings.database_url,
|
|
echo=False,
|
|
future=True,
|
|
pool_size=20, # Number of connections to maintain
|
|
max_overflow=10, # Additional connections when pool is exhausted
|
|
pool_pre_ping=True, # Verify connections before using
|
|
pool_recycle=3600, # Recycle connections after 1 hour
|
|
)
|
|
else:
|
|
# SQLite configuration - disable pooling to avoid locks
|
|
connect_args = {
|
|
"timeout": 30, # Wait up to 30 seconds for lock
|
|
"check_same_thread": False, # Allow usage across threads
|
|
}
|
|
|
|
engine = create_async_engine(
|
|
settings.database_url,
|
|
echo=False,
|
|
future=True,
|
|
connect_args=connect_args,
|
|
poolclass=NullPool, # Disable connection pooling for SQLite to avoid locks
|
|
pool_pre_ping=True, # Verify connections before using
|
|
)
|
|
|
|
AsyncSessionLocal = async_sessionmaker(
|
|
engine,
|
|
class_=AsyncSession,
|
|
expire_on_commit=False,
|
|
)
|
|
|
|
# Alias for compatibility
|
|
async_session_maker = AsyncSessionLocal
|
|
|
|
Base = declarative_base()
|
|
|
|
|
|
async def get_db():
|
|
async with AsyncSessionLocal() as session:
|
|
try:
|
|
yield session
|
|
finally:
|
|
await session.close()
|
|
|
|
|
|
async def init_db():
|
|
"""Initialize database connection pool
|
|
|
|
Note: Table creation is handled by Alembic migrations.
|
|
This function just ensures the connection pool is ready and optimizes settings.
|
|
"""
|
|
async with engine.begin() as conn:
|
|
# Detect database type from the actual engine URL, not just settings
|
|
# This is more reliable as it uses the actual connection
|
|
db_name = engine.url.drivername
|
|
|
|
if "sqlite" in db_name:
|
|
# SQLite-specific optimizations
|
|
await conn.execute(text("PRAGMA journal_mode=WAL"))
|
|
await conn.execute(text("PRAGMA synchronous=NORMAL"))
|
|
await conn.execute(text("PRAGMA busy_timeout=30000"))
|
|
elif "postgres" in db_name:
|
|
# PostgreSQL-specific optimizations (connection-level settings)
|
|
await conn.execute(text("SET timezone = 'UTC'"))
|
|
# Enable JIT compilation for complex queries (PG 11+)
|
|
await conn.execute(text("SET jit = on"))
|