56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
# utils/database.py
|
|
"""Database utilities for database operations.
|
|
|
|
This module provides utility functions and classes to interact with a database using SQLAlchemy. It includes:
|
|
- Creating a database engine with connection pooling.
|
|
- Generating a session factory for creating sessions.
|
|
"""
|
|
|
|
import logging
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker
|
|
from sqlalchemy.pool import QueuePool
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def get_db_engine(database_url: str):
|
|
"""Create a database engine with connection pooling.
|
|
|
|
Args:
|
|
database_url (str): The URL string to connect to the database. It can be for SQLite or PostgreSQL databases.
|
|
|
|
Returns:
|
|
sqlalchemy.engine.Engine: A SQLAlchemy Engine instance configured according to the provided database URL.
|
|
"""
|
|
if database_url.startswith("sqlite"):
|
|
# Configuration for SQLite database
|
|
engine = create_engine(
|
|
database_url, connect_args={"check_same_thread": False}, echo=False
|
|
)
|
|
else:
|
|
# Configuration for PostgreSQL databases with connection pooling
|
|
engine = create_engine(
|
|
database_url,
|
|
poolclass=QueuePool,
|
|
pool_size=10,
|
|
max_overflow=20,
|
|
pool_pre_ping=True,
|
|
echo=False,
|
|
)
|
|
|
|
logger.info(f"Database engine created for: " f"{database_url.split('@')[0]}@...")
|
|
return engine
|
|
|
|
|
|
def get_session_local(engine):
|
|
"""Create a session factory to generate database sessions.
|
|
|
|
Args:
|
|
engine (sqlalchemy.engine.Engine): The SQLAlchemy Engine instance created by `get_db_engine`.
|
|
|
|
Returns:
|
|
sqlalchemy.orm.sessionmaker: A session factory bound to the provided engine, suitable for creating database sessions.
|
|
"""
|
|
return sessionmaker(autocommit=False, autoflush=False, bind=engine)
|