base/alembic/env.py

0001 #!/usr/bin/env python
0002 #
0003 # Licensed under the Apache License, Version 2.0 (the "License");
0004 # You may not use this file except in compliance with the License.
0005 # You may obtain a copy of the License at
0006 # http://www.apache.org/licenses/LICENSE-2.0OA
0007 #
0008 # Authors:
0009 # - Wen Guan, <wen.guan@cern.ch>, 2023
0010
0011
0012 from logging.config import fileConfig
0013
0014 from sqlalchemy import engine_from_config
0015 from sqlalchemy import pool
0016
0017 from alembic import context
0018
0019 # from idds.orm.base.models import Base
0020 from idds.orm.base.session import BASE
0021
0022 # this is the Alembic Config object, which provides
0023 # access to the values within the .ini file in use.
0024 config = context.config
0025
0026 # Interpret the config file for Python logging.
0027 # This line sets up loggers basically.
0028 if config.config_file_name is not None:
0029     fileConfig(config.config_file_name)
0030
0031 # add your model's MetaData object here
0032 # for 'autogenerate' support
0033 # from myapp import mymodel
0034 # target_metadata = mymodel.Base.metadata
0035 # target_metadata = None
0036 target_metadata = BASE.metadata
0037
0038 # other values from the config, defined by the needs of env.py,
0039 # can be acquired:
0040 # my_important_option = config.get_main_option("my_important_option")
0041 # ... etc.
0042
0043
0044 def run_migrations_offline() -> None:
0045     """Run migrations in 'offline' mode.
0046
0047     This configures the context with just a URL
0048     and not an Engine, though an Engine is acceptable
0049     here as well.  By skipping the Engine creation
0050     we don't even need a DBAPI to be available.
0051
0052     Calls to context.execute() here emit the given string to the
0053     script output.
0054
0055     """
0056     url = config.get_main_option("sqlalchemy.url")
0057
0058     version_table_schema = config.get_main_option("version_table_schema")
0059
0060     context.configure(
0061         url=url,
0062         version_table_schema=version_table_schema,
0063         target_metadata=target_metadata,
0064         literal_binds=True,
0065         dialect_opts={"paramstyle": "named"},
0066         include_schemas=True
0067     )
0068
0069     with context.begin_transaction():
0070         context.run_migrations()
0071
0072
0073 def run_migrations_online() -> None:
0074     """Run migrations in 'online' mode.
0075
0076     In this scenario we need to create an Engine
0077     and associate a connection with the context.
0078
0079     """
0080     connectable = engine_from_config(
0081         config.get_section(config.config_ini_section),
0082         prefix="sqlalchemy.",
0083         poolclass=pool.NullPool,
0084     )
0085
0086     with connectable.connect() as connection:
0087         context.configure(
0088             connection=connection, target_metadata=target_metadata,
0089             version_table_schema=config.get_main_option("version_table_schema"),
0090             include_schemas=True
0091         )
0092
0093         with context.begin_transaction():
0094             context.run_migrations()
0095
0096
0097 if context.is_offline_mode():
0098     run_migrations_offline()
0099 else:
0100     run_migrations_online()