env.py 3.58 KB
Newer Older
Lukáš Lalinský's avatar
Lukáš Lalinský committed
1 2
from __future__ import with_statement
import os
Lukáš Lalinský's avatar
Lukáš Lalinský committed
3
import logging
Lukáš Lalinský's avatar
Lukáš Lalinský committed
4 5 6 7 8 9
from alembic import context
from sqlalchemy import engine_from_config, pool
from logging.config import fileConfig

config = context.config
fileConfig(config.config_file_name)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
10
logger = logging.getLogger("alembic.env")
Lukáš Lalinský's avatar
Lukáš Lalinský committed
11 12 13 14 15

import acoustid.tables
target_metadata = acoustid.tables.metadata

import acoustid.config
16 17
acoustid_config_filename = os.environ.get('ACOUSTID_CONFIG',
    os.path.join(os.path.dirname(__file__), '..', 'acoustid.conf'))
18 19 20
acoustid_config = acoustid.config.Config()
acoustid_config.read(acoustid_config_filename)
acoustid_config.read_env()
Lukáš Lalinský's avatar
Lukáš Lalinský committed
21

Lukáš Lalinský's avatar
Lukáš Lalinský committed
22 23 24 25 26 27 28 29
use_two_phase_commit = acoustid_config.databases.use_two_phase_commit


def include_object(db_name):
    def inner(obj, name, obj_type, reflected, compare_to):
        if obj_type == "table":
            if obj.schema == "musicbrainz":
                return False
30
            bind_key = obj.info.get('bind_key', 'main')
Lukáš Lalinský's avatar
Lukáš Lalinský committed
31 32 33 34 35
            if bind_key != db_name:
                return False
        if obj_type == "column":
            if obj.table.schema == "musicbrainz":
                return False
36
            bind_key = obj.table.info.get('bind_key', 'main')
Lukáš Lalinský's avatar
Lukáš Lalinský committed
37 38 39 40
            if bind_key != db_name:
                return False
        return True
    return inner
Lukáš Lalinský's avatar
Lukáš Lalinský committed
41 42 43 44 45 46 47 48 49 50 51 52 53 54


def run_migrations_offline():
    """Run migrations in 'offline' mode.

    This configures the context with just a URL
    and not an Engine, though an Engine is acceptable
    here as well.  By skipping the Engine creation
    we don't even need a DBAPI to be available.

    Calls to context.execute() here emit the given string to the
    script output.

    """
Lukáš Lalinský's avatar
Lukáš Lalinský committed
55 56 57 58 59 60 61 62 63
    for name, db_config in acoustid_config.databases.databases.items():
        logger.info("Migrating database %s" % name)

        context.configure(
            url=db_config.create_url(),
            target_metadata=target_metadata,
            literal_binds=True,
            include_object=include_object(name),
        )
Lukáš Lalinský's avatar
Lukáš Lalinský committed
64

Lukáš Lalinský's avatar
Lukáš Lalinský committed
65
        with context.begin_transaction():
Lukáš Lalinský's avatar
Lukáš Lalinský committed
66
            context.run_migrations(engine_name=name)
Lukáš Lalinský's avatar
Lukáš Lalinský committed
67 68 69 70 71 72 73 74 75


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.

    """
Lukáš Lalinský's avatar
Lukáš Lalinský committed
76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114
    engines = {}
    for name, db_config in acoustid_config.databases.databases.items():
        engines[name] = rec = {}
        rec["engine"] = db_config.create_engine(poolclass=pool.NullPool)

    for name, rec in engines.items():
        engine = rec["engine"]
        rec["connection"] = conn = engine.connect()

        if use_two_phase_commit:
            rec["transaction"] = conn.begin_twophase()
        else:
            rec["transaction"] = conn.begin()

    try:
        for name, rec in engines.items():
            logger.info("Migrating database %s" % name)
            context.configure(
                connection=rec["connection"],
                upgrade_token="%s_upgrades" % name,
                downgrade_token="%s_downgrades" % name,
                target_metadata=target_metadata,
                include_object=include_object(name),
            )
            context.run_migrations(engine_name=name)

        if use_two_phase_commit:
            for rec in engines.values():
                rec["transaction"].prepare()

        for rec in engines.values():
            rec["transaction"].commit()
    except:
        for rec in engines.values():
            rec["transaction"].rollback()
        raise
    finally:
        for rec in engines.values():
            rec["connection"].close()
Lukáš Lalinský's avatar
Lukáš Lalinský committed
115 116 117 118 119 120


if context.is_offline_mode():
    run_migrations_offline()
else:
    run_migrations_online()