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.")