import logging
from typing import Tuple, Optional
from sqlalchemy import create_engine, Engine, text, event
from sqlalchemy.orm import declarative_base, scoped_session, sessionmaker
from sqlalchemy.pool import NullPool
from urllib.parse import quote_plus
from .config.config import ccat_ops_db_settings
# Add this near the top of the file
logger = logging.getLogger(__name__)
logger.propagate = True
Base = declarative_base()
# Cache engines by (database_type, host, port, db_name) to avoid recreating
# engine + running create_all() on every request.
_engine_cache: dict[tuple, Engine] = {}
# Guards create_schema() so create_all runs at most once per engine. Keyed by
# the Engine object (one engine per cache_key), not the tuple, so the guard and
# the thing it guards share identity.
_schema_initialized: set = set()
[docs]
def get_database_url(
database_type: str,
database: Optional[str] = None,
host: Optional[str] = None,
port: Optional[str] = None,
user: Optional[str] = None,
password: Optional[str] = None,
db_name: Optional[str] = None,
async_driver: bool = False,
) -> str:
"""
Generate database URL based on type and configuration.
Args:
database_type: Type of database ('sqlite', 'mysql', 'postgresql')
database: Optional database URL override
host: Optional host override
port: Optional port override
user: Optional user override
password: Optional password override
db_name: Optional database name override
async_driver: Whether to use async driver (e.g., asyncpg for PostgreSQL)
"""
if database_type == "sqlite":
database = database or ccat_ops_db_settings.DATABASE_SQLITE_DATABASE
return f"sqlite:///{database}"
elif database_type == "mysql":
user = user or ccat_ops_db_settings.DATABASE_MYSQL_USER
password = password or ccat_ops_db_settings.DATABASE_MYSQL_PASSWORD
host = host or ccat_ops_db_settings.DATABASE_MYSQL_HOST
port = port or ccat_ops_db_settings.DATABASE_MYSQL_PORT
db = db_name or database or ccat_ops_db_settings.DATABASE_MYSQL_DATABASE
return f"mysql+mysqldb://{user}:{quote_plus(password)}@{host}:{port}/{db}"
elif database_type == "postgresql":
user = user or ccat_ops_db_settings.DATABASE_POSTGRESQL_USER
password = password or ccat_ops_db_settings.DATABASE_POSTGRESQL_PASSWORD
host = host or ccat_ops_db_settings.DATABASE_POSTGRESQL_HOST
port = port or ccat_ops_db_settings.DATABASE_POSTGRESQL_PORT
db = db_name or database or ccat_ops_db_settings.DATABASE_POSTGRESQL_DATABASE
if async_driver:
return (
f"postgresql+asyncpg://{user}:{quote_plus(password)}@{host}:{port}/{db}"
)
else:
return f"postgresql://{user}:{quote_plus(password)}@{host}:{port}/{db}"
else:
raise ValueError(f"Unsupported database type: {database_type}")
def get_engine(
database_type: Optional[str] = None,
database: Optional[str] = None,
null_pool: bool = False,
host: Optional[str] = None,
port: Optional[str] = None,
user: Optional[str] = None,
password: Optional[str] = None,
db_name: Optional[str] = None,
) -> Engine:
"""Create (or return a cached) SQLAlchemy engine. Runs no DDL.
Engine/connection creation is a runtime operation; schema management
(``create_schema``) is a one-time deployment operation. Keeping them in
separate functions (ops-db #58 item 4) means a service that only needs a
session never triggers ``create_all`` as a side effect.
Engines are cached by (database_type, host, port, db_name, null_pool) so
repeated calls reuse the connection pool.
"""
database_type = database_type or ccat_ops_db_settings.DATABASE_TYPE
# Resolve actual values for cache key (before engine creation)
_host = host or ccat_ops_db_settings.DATABASE_POSTGRESQL_HOST
_port = port or ccat_ops_db_settings.DATABASE_POSTGRESQL_PORT
_db_name = db_name or database
cache_key = (database_type, _host, _port, _db_name, null_pool)
if cache_key in _engine_cache:
return _engine_cache[cache_key]
logger.info(
"Using database_type %s host %s port %s",
database_type,
_host,
_port,
)
url = get_database_url(database_type, database, host, port, user, password, db_name)
engine_kwargs = {
"echo": False,
"pool_pre_ping": True,
}
if null_pool:
engine_kwargs["poolclass"] = NullPool
engine = create_engine(url, **engine_kwargs)
if database_type == "postgresql":
@event.listens_for(engine, "connect")
def set_timezone(dbapi_conn, connection_record):
cursor = dbapi_conn.cursor()
cursor.execute("SET timezone='UTC'")
cursor.close()
_engine_cache[cache_key] = engine
return engine
def create_schema(engine: Engine, drop: bool = False) -> None:
"""Create the ORM schema on ``engine`` via ``Base.metadata.create_all``.
Separated from session/engine creation (ops-db #58 item 4) so it is the
only place ``create_all`` runs. NOTE (ops-db #77): production-style
provisioning (opsdb_init) no longer calls this — it builds the schema with
Alembic migrations (``opsdb_migrate``) so staging mirrors the production
migrate path. ``create_all`` here remains for the legacy/local/sqlite
bootstrap (the non ``--seed-only`` opsdb_init path and direct dev use).
create_all runs at most once per engine (guarded by ``_schema_initialized``)
unless ``drop=True``, which always drops and rebuilds.
"""
is_postgres = engine.dialect.name == "postgresql"
if drop:
logger.info("Dropping all tables and types")
# For PostgreSQL, we need to drop enum types first
if is_postgres:
with engine.connect() as conn:
conn.execute(text("DROP TYPE IF EXISTS status CASCADE"))
conn.commit()
Base.metadata.drop_all(bind=engine)
Base.metadata.create_all(bind=engine)
_schema_initialized.add(engine)
return
if engine in _schema_initialized:
return
if is_postgres:
# For PostgreSQL, we need to ensure the enum type exists
with engine.connect() as conn:
# Check if the enum type exists
result = conn.execute(
text("SELECT 1 FROM pg_type WHERE typname = 'status'")
)
if not result.scalar():
# Create the enum type if it doesn't exist
conn.execute(
text(
"""
CREATE TYPE status AS ENUM (
'pending',
'scheduled',
'in_progress',
'completed',
'failed'
)
"""
)
)
conn.commit()
Base.metadata.create_all(bind=engine)
_schema_initialized.add(engine)
def get_session_factory(engine: Engine) -> scoped_session:
"""Return a ``scoped_session`` bound to ``engine``.
Kept as ``scoped_session`` (not a bare ``sessionmaker``) for backward
compatibility: data-transfer and workflow-manager rely on
``scoped_session.remove()`` for connection cleanup. Migrating callers to a
plain ``sessionmaker`` is ops-db #58 item 7, deliberately deferred to a
separate coordinated cross-repo change.
"""
return scoped_session(sessionmaker(bind=engine))
[docs]
def init_ccat_ops_db(
database_type: Optional[str] = None,
database: Optional[str] = None,
drop: bool = False,
null_pool: bool = False,
host: Optional[str] = None,
port: Optional[str] = None,
user: Optional[str] = None,
password: Optional[str] = None,
db_name: Optional[str] = None,
read_only: bool = False,
) -> Tuple[scoped_session, Engine]:
"""
Initialize a database connection based on the supplied configuration.
Backward-compatible composition over ``get_engine`` / ``create_schema`` /
``get_session_factory`` (ops-db #58). Behavior and return type are unchanged
for existing consumers (ops-db-api, data-transfer, workflow-manager): it
returns ``(scoped_session, engine)`` and, unless ``read_only=True``, ensures
the schema exists via ``create_all``.
Parameters:
-----------
database_type: str, optional
Can be 'sqlite', 'mysql', or 'postgresql'. Defaults to config setting.
database: str, optional
URL string to connect to the database. Defaults to config setting.
drop: bool, default False
If True, drops all tables before creating them.
null_pool: bool, default False
If True, uses NullPool instead of the default connection pool.
host: str, optional
Database host override.
port: str, optional
Database port override.
user: str, optional
Database user override.
password: str, optional
Database password override.
db_name: str, optional
Database name override.
read_only: bool, default False
If True, only opens the connection — never runs DDL.
Returns:
--------
Tuple[scoped_session, Engine]
A tuple containing the database session and engine.
"""
engine = get_engine(
database_type=database_type,
database=database,
null_pool=null_pool,
host=host,
port=port,
user=user,
password=password,
db_name=db_name,
)
if not read_only:
create_schema(engine, drop=drop)
return get_session_factory(engine), engine
