def db() -> Dict[Engine, sessionmaker]:
engine = create_engine(TestConfig.SQLALCHEMY_DATABASE_URI, echo=True)
session = sessionmaker(bind=engine)
_db = {'engine': engine, 'session': session}
try:
alembic_config = AlembicConfig(os.path.abspath("../../alembic.ini"))
alembic_config.set_main_option('script_location',
os.path.abspath("../../meant_alembic"))
alembic_config.set_main_option('sqlalchemy.url',
TestConfig.SQLALCHEMY_DATABASE_URI)
alembic_upgrade(alembic_config, 'head')
except CommandError:
log(message="testing only specified TCs", keyword="INFO")
alembic_config = AlembicConfig(os.path.abspath("../../../alembic.ini"))
alembic_config.set_main_option(
'script_location', os.path.abspath("../../../meant_alembic"))
alembic_config.set_main_option('sqlalchemy.url',
TestConfig.SQLALCHEMY_DATABASE_URI)
alembic_upgrade(alembic_config, 'head')
log(message="database created", keyword="INFO")
yield _db
log(message="database disposed", keyword="INFO")
engine.dispose()
def _get_alembic_config_from_cache(
force_cfg: Optional[Dict] = None, ) -> Optional[AlembicConfig]:
"""
Creates alembic config from cfg or cache
Returns None if cannot build url (e.g. if user requires a cache that does not exists)
"""
# build url
try:
if force_cfg:
cfg = force_cfg
else:
cfg = _load_cache(raise_if_error=True)
url = build_url(**cfg)
except Exception:
log.debug("Cannot open cache or cannot build URL",
exc_info=True,
stack_info=True)
click.echo("Invalid database config, please run discover first",
err=True)
_reset_cache()
return None
# build config
config = AlembicConfig(default_ini)
config.set_main_option("script_location", str(migration_dir))
config.set_main_option("sqlalchemy.url", str(url))
return config
def _alembic_config(connection):
root = t.pipe(__file__, os.path.realpath, os.path.dirname)
config = AlembicConfig(os.path.join(root, 'alembic.ini'))
config.set_main_option('script_location',
os.path.join(root, 'migrations'))
config.attributes['connection'] = connection
return config
def db(app, request):
"""db
Returns session-wide initialized database.
"""
database_uri = sqlalchemy.engine.url.make_url(
TestConfig.SQLALCHEMY_DATABASE_URI)
host_uri = sqlalchemy.engine.url.URL(database_uri.drivername,
username=database_uri.username,
password=database_uri.password,
host=database_uri.host,
port=database_uri.port)
database_name = database_uri.database
template_engine = sqlalchemy.create_engine(host_uri, echo=False)
conn = template_engine.connect()
conn = conn.execution_options(autocommit=True,
isolation_level='AUTOCOMMIT')
try:
conn.execute(f'DROP DATABASE IF EXISTS {database_name};')
conn.execute(f'CREATE DATABASE {database_name};')
except:
pass
finally:
conn.close()
template_engine.dispose()
Migrate(_db.app, _db)
alembic_config = AlembicConfig('/srv/migrations/alembic.ini')
alembic_config.set_main_option('script_location', 'migrations')
command.upgrade(alembic_config, 'head')
yield _db
def get_alembic_config():
"""Returns an alembic Config with the right ini file."""
# alembic.ini is in the repository root.
alembic_ini = str(
pathlib.Path(__file__).parent.parent.parent / "alembic.ini")
assert os.path.exists(alembic_ini)
return AlembicConfig(alembic_ini)
def perform_migratons(config_name):
''' If fails, then we should revert to previous version of SLUG running on Heroku
link: http://stackoverflow.com/questions/24622170/using-alembic-api-from-inside-application-code
'''
db_url = configuration[config_name].SQLALCHEMY_DATABASE_URI
alembic_config = AlembicConfig('.\\AlertWeb\\alembic.ini')
alembic_config.set_main_option('sqlalchemy.url', db_url)
alembic_config.set_main_option('script_location',
'.\\AlertWeb\\migrations')
script_dir = ScriptDirectory.from_config(alembic_config)
head_revision = script_dir.get_current_head()
current_revision = get_current_revision(db_url)
def upgrade(rev, context):
print(rev)
return script_dir._upgrade_revs(head_revision, rev)
#script_dir.
# Facade for migration context.
with EnvironmentContext(alembic_config,
script_dir,
as_sql=False,
fn=upgrade,
starting_rev=current_revision,
destination_rev=head_revision,
tag=None):
script_dir.run_env()
def list_migrations(cfg_path, head):
cfg = AlembicConfig(cfg_path)
script = ScriptDirectory.from_config(cfg)
migrations = [x.revision
for x in script.walk_revisions(base='base', head=head)]
migrations.reverse()
return migrations
def pg_server(unused_port, container_starter):
tag = "13.1"
image = 'postgres:{}'.format(tag)
internal_port = 5432
host_port = unused_port()
environment = {'POSTGRES_USER': '******',
'POSTGRES_PASSWORD': '******',
'POSTGRES_DB': 'postgres'}
#
# volume = (str(TEMP_FOLDER / 'docker' / 'pg'),
# '/var/lib/postgresql/data')
container = container_starter(image, 'db', internal_port, host_port,
environment, None)
params = dict(database='postgres',
user='******',
password='******',
host='localhost',
port=host_port)
def connect():
conn = psycopg2.connect(**params)
cur = conn.cursor()
cur.close()
conn.close()
wait_for_container(connect, image, psycopg2.Error)
container['params'] = params
alembic_config = AlembicConfig('alembic.ini')
alembic_config.set_main_option('sqlalchemy.url', f'postgresql://*****:*****@localhost:{container["port"]}/postgres')
alembic_upgrade(alembic_config, 'head')
return container
def setup_migration_version_control(self):
self.reset_alembic_output()
alembic_config = AlembicConfig()
alembic_config.set_main_option(
"script_location", self.migrate_repository
)
alembic_config.set_main_option(
"sqlalchemy.url", str(self.metadata.bind.url)
)
try:
sqlalchemy_migrate_version = self.metadata.bind.execute(
u'select version from migrate_version'
).scalar()
except ProgrammingError:
sqlalchemy_migrate_version = 0
# this value is used for graceful upgrade from
# sqlalchemy-migrate to alembic
alembic_config.set_main_option(
"sqlalchemy_migrate_version", str(sqlalchemy_migrate_version)
)
# This is an interceptor for alembic output. Otherwise,
# everything will be printed to stdout
alembic_config.print_stdout = self.add_alembic_output
self.alembic_config = alembic_config
def init_db(environment="default"):
SQLALCHEMY_DATABASE_URL = config[environment]().DB_CONNECTION
# SQLALCHEMY_DATABASE_URL = "sqlite:///./sql_app.db"
engine = create_engine(
SQLALCHEMY_DATABASE_URL
)
# TODO: Soe schools of thought on DB migrations say run the create_all first for setting up an up-to-date local
# environment in one command instead of iterating through all the migration versions. Laving this here
# incase I want to do that. Personally, I'd like to treat local dev like prod, where create_all will never run.
# models.Base.metadata.create_all(bind=engine)
alembic_cfg = AlembicConfig("./app/alembic.ini")
alembic_cfg.set_main_option("sqlalchemy.url", config[environment]().DB_CONNECTION)
command.upgrade(alembic_cfg, "head")
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
global session
session = SessionLocal
# TODO: These return values aren't being used, but I'd like to use them over fastapi's dependency injection pattern
# as a way to provide the DB connection
return session, engine
def version_schema_new(script_location: str):
"""Applies alembic versioning to schema."""
# add it to alembic table
alembic_cfg = AlembicConfig(config.ALEMBIC_INI_PATH)
alembic_cfg.set_main_option("script_location", script_location)
alembic_command.stamp(alembic_cfg, "head")
def upgrade_database(tag, sql, revision):
"""Upgrades database schema to newest version."""
alembic_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
"alembic.ini")
alembic_cfg = AlembicConfig(alembic_path)
alembic_command.upgrade(alembic_cfg, revision, sql=sql, tag=tag)
click.secho("Success.", fg="green")
def db(request):
host_url = f'postgres+psycopg2://{POSTGRES_USER}:{POSTGRES_PASS}@{POSTGRES_HOST}:{POSTGRES_PORT}'
pg_database_url = '{}/postgres'.format(host_url)
engine = create_engine(pg_database_url, echo=True)
conn = engine.connect()
conn.execute('commit')
try:
conn.execute('create database test')
except:
pass
conn.close()
alembic_config = AlembicConfig('alembic.ini')
alembic_config.set_main_option('sqlalchemy.url',
'{}/test'.format(host_url))
alembic_upgrade(alembic_config, 'head')
yield
conn = engine.connect()
conn.execute(
"SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = 'test';"
)
conn.execute('commit')
conn.execute('DROP DATABASE test')
conn.close()
engine.dispose()
def upgrade_database(tag, sql, revision):
"""Upgrades database schema to newest version."""
from sqlalchemy_utils import database_exists, create_database
from alembic.migration import MigrationContext
org_engine = engine
ORG_SQLALCHEMY_DATABASE_URI = config.SQLALCHEMY_DATABASE_URI
# print("config.SQLALCHEMY_DATABASE_URI=", config.SQLALCHEMY_DATABASE_URI)
alembic_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
"alembic.ini")
alembic_cfg = AlembicConfig(alembic_path)
if not database_exists(str(ORG_SQLALCHEMY_DATABASE_URI)): #
create_database(str(ORG_SQLALCHEMY_DATABASE_URI))
Base.metadata.create_all(org_engine)
alembic_command.stamp(alembic_cfg, "head")
else:
conn = org_engine.connect()
context = MigrationContext.configure(conn)
current_rev = context.get_current_revision()
if not current_rev:
Base.metadata.create_all(org_engine)
alembic_command.stamp(alembic_cfg, "head")
else:
alembic_command.upgrade(alembic_cfg, revision, sql=sql, tag=tag)
sync_triggers()
click.secho("Success. The database is upgraded.", fg="green")
def _alembic_config(self):
dir = os.path.join(os.path.dirname(__file__), 'migrations')
config = AlembicConfig(os.path.join(dir, 'alembic.ini'))
config.set_main_option('script_location', dir)
config.set_main_option('sqlalchemy.url', self.db_uri)
config.set_main_option('adnotatio_server_path',
os.path.dirname(__file__))
return config
def reset_db():
"""Resets the database to the original state using alembic downgrade and upgrade commands"""
from alembic.command import downgrade, upgrade
from alembic.config import Config as AlembicConfig
config = AlembicConfig('alembic.ini')
downgrade(config, 'base')
upgrade(config, 'head')
click.echo('Database has been reset')
def get_config():
path = os.path.join(os.path.abspath(os.path.dirname(__file__)),
"migrations")
config = AlembicConfig(os.path.join(path, "alembic.ini"))
config.set_main_option("script_location", path)
config.cmd_opts = argparse.Namespace()
setattr(config.cmd_opts, "x", None)
return config
def __init__(self, config: 'Config'):
self._alembic_config = AlembicConfig()
self._alembic_config.set_main_option(
"script_location",
"{{cookiecutter.project_name}}/internal/alembic")
self._alembic_config.set_main_option("sqlalchemy.url", config.database)
self._migration_message = config.migration_creator.migration_message
self._auto_generate = config.migration_creator.auto_generate
def build_alembic_config(postgresql_uri):
alembic_path = "{}/alembic".format(config.get('repos', 'manage_db'))
ini_file = "{}/alembic.ini".format(config.get('repos', 'manage_db'))
alembic_cfg = AlembicConfig(ini_file)
alembic_cfg.set_main_option('configure_logging', 'false')
alembic_cfg.set_main_option("script_location", alembic_path)
alembic_cfg.set_main_option("sqlalchemy.url", postgresql_uri)
return alembic_cfg
def db_migrations(app):
alembic_ini = os.path.join(config.project_dir, 'alembic.ini')
alembic_config = AlembicConfig(alembic_ini)
alembic_config.attributes['testing'] = True
upgrade(alembic_config, 'head')
yield
database.db.session.remove()
database.db.drop_all(app=app)
stamp(alembic_config, 'base')
def __init__(self, config: 'Config'):
self._alembic_config = AlembicConfig()
self._alembic_config.set_main_option(
"script_location",
"{{cookiecutter.project_name}}/internal/alembic")
self._alembic_config.set_main_option("sqlalchemy.url", config.database)
self._revision = config.migrator.revision
self._sql_only = config.migrator.sql_only
def is_alembic_head():
alembic_cfg_path = os.path.join(os.path.dirname(os.path.abspath(__file__)),
'alembic.ini')
alembic_cfg = AlembicConfig(alembic_cfg_path)
context = MigrationContext.configure(db_session.connection())
script = ScriptDirectory.from_config(alembic_cfg)
current_revision = context.get_current_revision()
head_revision = script.get_current_head()
return current_revision == head_revision
def upgrade_database(tag, sql, revision, revision_type):
"""Upgrades database schema to newest version."""
import sqlalchemy
from sqlalchemy import inspect
from sqlalchemy_utils import database_exists
from alembic import command as alembic_command
from alembic.config import Config as AlembicConfig
from .database import engine
from .database_util.manage import (
get_core_tables,
get_tenant_tables,
init_database,
setup_fulltext_search,
)
alembic_cfg = AlembicConfig(config.ALEMBIC_INI_PATH)
if not database_exists(str(config.SQLALCHEMY_DATABASE_URI)):
click.secho("Found no database to upgrade, initializing new database...")
init_database(engine)
else:
conn = engine.connect()
# detect if we need to convert to a multi-tenant schema structure
schema_names = inspect(engine).get_schema_names()
if "dispatch_core" not in schema_names:
click.secho("Detected single tenant database, converting to multi-tenant...")
conn.execute(sqlalchemy.text(open(config.ALEMBIC_MULTI_TENANT_MIGRATION_PATH).read()))
# init initial triggers
conn.execute("set search_path to dispatch_core")
setup_fulltext_search(conn, get_core_tables())
tenant_tables = get_tenant_tables()
for t in tenant_tables:
t.schema = "dispatch_organization_default"
conn.execute("set search_path to dispatch_organization_default")
setup_fulltext_search(conn, tenant_tables)
if revision_type:
if revision_type == "core":
path = config.ALEMBIC_CORE_REVISION_PATH
elif revision_type == "tenant":
path = config.ALEMBIC_TENANT_REVISION_PATH
alembic_cfg.set_main_option("script_location", path)
alembic_command.upgrade(alembic_cfg, revision, sql=sql, tag=tag)
else:
for path in [config.ALEMBIC_CORE_REVISION_PATH, config.ALEMBIC_TENANT_REVISION_PATH]:
alembic_cfg.set_main_option("script_location", path)
alembic_command.upgrade(alembic_cfg, revision, sql=sql, tag=tag)
click.secho("Success.", fg="green")
def _db(app):
db = SQLAlchemy(app)
with app.app_context():
Migrate(app, db)
alembic_config = AlembicConfig('migrations/alembic.ini')
alembic_config.set_main_option('sqlalchemy.url',
app.config['SQLALCHEMY_DATABASE_URI'])
alembic_upgrade(alembic_config, 'head')
return db
def get_db_config():
from alembic.config import Config as AlembicConfig
INI_FILE = os.path.abspath(
os.path.join(os.path.dirname(__file__), '..', 'alembic.ini'))
alembic_oonfig = AlembicConfig(INI_FILE)
alembic_oonfig.set_main_option('sqlalchemy.url',
app.config['SQLALCHEMY_DATABASE_URI'])
return alembic_oonfig
def migrate():
alembic_config = AlembicConfig(
os.path.join(tuber.__path__[0], "alembic.ini"))
if oneshot_db_create:
# To avoid running migrations on sqlite dev databases just create the current
# tables and stamp them as being up to date so that migrations won't run.
# This should only run if there is not an existing db.
create_tables()
alembic.command.stamp(alembic_config, "head")
alembic.command.upgrade(alembic_config, "head")
def downgrade_database(tag, sql, revision):
"""Downgrades database schema to next newest version."""
alembic_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "alembic.ini")
alembic_cfg = AlembicConfig(alembic_path)
if sql and revision == "-1":
revision = "head:-1"
alembic_command.downgrade(alembic_cfg, revision, sql=sql, tag=tag)
click.secho("Success.", fg="green")
def create_alembic_config(**kwargs):
"""Returns an `alembic.config.Config` object configured for uber.
"""
kwargs['file_'] = alembic_ini_path
alembic_config = AlembicConfig(**kwargs)
# Override settings from "alembic.ini"
alembic_config.set_main_option('script_location', script_location)
alembic_config.set_main_option(
'version_locations', version_locations_option)
return alembic_config
def _create_database():
"""
Creates/Updates the database.
"""
directory = os.path.join(os.path.dirname(__file__), '../migrations')
config = AlembicConfig(os.path.join(
directory,
'alembic.ini'
))
config.set_main_option('script_location', directory)
command.upgrade(config, 'head', sql=False, tag=None)
def build_alembic_config(engine: Engine) -> AlembicConfig:
"""Populate alembic configuration from metadata and config file."""
path_to_alembic_ini = REPO_ROOT / "alembic.ini"
alembic_cfg = AlembicConfig(path_to_alembic_ini)
# Make double sure alembic references the test database
# print(" UPGRADE-----------------"+str(engine.url))
alembic_cfg.set_main_option("sqlalchemy.url", str(engine.url))
alembic_cfg.set_main_option("script_location",
str((Path("src") / "test" / "alembic_config")))
return alembic_cfg
