def clean_test_db_engine(use_in_memory_db):
if use_in_memory_db:
in_memory_database_url = "sqlite+pysqlite:///:memory:"
engine = get_db_engine(override_db_url=in_memory_database_url)
else:
engine = get_db_engine()
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)
return engine
def detect_in_memory_db() -> bool:
from hetdesrun.persistence import get_db_engine
engine = get_db_engine()
backend_name = engine.url.get_backend_name()
# driver_name = engine.url.get_driver_name() # pysqlite
database = engine.url.database # ":memory:"
if backend_name.lower() == "sqlite" and (
(database is None) or database.lower() in (":memory:", )):
return True
return False
def run_migrations_online():
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine
and associate a connection with the context.
"""
from hetdesrun import configure_logging
logger = logging.getLogger("hetdesrun")
alembic_logger = logging.getLogger("alembic")
configure_logging(alembic_logger)
from hetdesrun.persistence import get_db_engine
connectable = get_db_engine()
logger.info("Running online migrations with driver %s",
connectable.url.get_driver_name())
with connectable.connect() as connection:
context.configure(connection=connection,
target_metadata=get_target_metadata())
logger.info("Connected to db for migrations")
with context.begin_transaction():
logger.info("Beginning transaction. Dialect name is %s",
connection.dialect.name)
if connection.dialect.name == "postgresql":
logger.info(
"Detected postgresql driver. Ensuring versioning table")
# Make sure no two processed can migrate at the same time
context.get_context()._ensure_version_table() # pylint: disable=protected-access
logger.info(
"Ensured versioning table. Now locking alembic version table"
)
connection.execute(
text(
"LOCK TABLE alembic_version IN ACCESS EXCLUSIVE MODE"))
# Postgres lock is released when transaction ends
logger.info("actually starting to run migrations")
context.run_migrations()
def run_migrations(
alembic_dir: str = "./alembic",
connection_url=get_config().sqlalchemy_connection_string,
) -> None:
"""Runs alembic migrations from within Python code
Should only be used for local development server. Not recommended
for multi-process/thread production servers.
Note: The docker container runs migrations via prestart.sh script in the
production setup.
"""
from hetdesrun import migrations_invoked_from_py
migrations_invoked_from_py = True
from pydantic import SecretStr
import hetdesrun.persistence.dbmodels
from alembic import command
from alembic.config import Config
from hetdesrun.persistence import get_db_engine
engine = get_db_engine()
logger.info("Using DB engine driver: %s", str(engine.url.drivername))
if isinstance(connection_url, SecretStr):
connection_url_to_use = connection_url.get_secret_value()
else:
connection_url_to_use = connection_url
logger.info("Running DB migrations in %s", alembic_dir)
alembic_cfg = Config()
alembic_cfg.set_main_option("script_location", alembic_dir)
# alembic_cfg.set_main_option("sqlalchemy.url", connection_url_to_use)
# alembic_cfg.set_section_option("logger_root", "level", "DEBUG")
# alembic_cfg.set_section_option("logger_alembic", "level", "DEBUG")
# alembic_cfg.set_section_option("logger_sqlalchemy", "level", "DEBUG")
command.upgrade(alembic_cfg, "head")
logger.info("Finished running migrations.")
