Passed
Pull Request — dev (#82)
by Stephan
01:10
created

env.run_migrations_offline()   A

Complexity

Conditions 2

Size

Total Lines 22
Code Lines 9

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 2
eloc 9
nop 0
dl 0
loc 22
rs 9.95
c 0
b 0
f 0
1
from fileinput import FileInput
2
from logging.config import fileConfig
3
from pathlib import Path
4
import re
5
6
from alembic import context
7
from alembic.script import ScriptDirectory, write_hooks
8
9
from egon.data.db import engine
10
from egon.data.orm.openstreetmap import metadata
11
12
# this is the Alembic Config object, which provides
13
# access to the values within the .ini file in use.
14
config = context.config
15
16
# Interpret the config file for Python logging.
17
# This line sets up loggers basically.
18
fileConfig(config.config_file_name)
19
20
# add your model's MetaData object here
21
# for 'autogenerate' support
22
# from myapp import mymodel
23
# target_metadata = mymodel.Base.metadata
24
target_metadata = metadata
25
26
# other values from the config, defined by the needs of env.py,
27
# can be acquired:
28
# my_important_option = config.get_main_option("my_important_option")
29
# ... etc.
30
31
32
@write_hooks.register("fixup")
33
def fix_minor_issues(filename, options):
34
    """Fix minor issues like grammar in comments or trailing whitespace."""
35
    fixes = []
36
    with FileInput(filename, inplace=True) as script:
37
        for line in script:
38
            if re.match(r"# revision identifiers, used by Alembic\.", line):
39
                print("# Revision identifiers, used by Alembic.")
40
                fixes.append('"# revision ..." -> "# Revision ..."')
41
            else:
42
                print(line, end="")
43
    for fix in fixes:
44
        print(f"Fixed: {fix}.")
45
46
47
def script(filename=None, revision=None):
48
    """Return the script object matching the parameters.
49
50
    Either the script object having the given :py:obj:`filename` or having
51
    the given :py:obj:`revision`, is returned. If both are given, both have
52
    to match.
53
54
    Supplying at least one is mandatory.
55
    """
56
    assert filename is not None or revision is not None
57
    scripts = ScriptDirectory.from_config(config)
58
    return [
59
        s
60
        for s in scripts.walk_revisions()
61
        if (filename is None or s.path == filename)
62
        and (revision is None or s.revision == revision)
63
    ][0]
64
65
66
@write_hooks.register("sequentialize")
67
def assign_sequential_version_number(filename, options):
68
    """Prefix the migration script's name with a sequential version number.
69
70
    The version number is generated by incrementing the version number of
71
    the down revision by one. Since the intention is to get a sequential
72
    ordering when sorting the migration scripts alphabetically, this
73
    sequential number is right aligned and padded with zeroes to a length of
74
    three characters.
75
    """
76
77
    current = script(filename=filename)
78
    path = Path(filename)
79
    down = (
80
        "-001."  # Gets incremented to 0 in the line below.
81
        if current.down_revision is None
82
        else Path(script(revision=current.down_revision).path).name
83
    )
84
    version = int(re.match(r"^(-?\d{3})\..*", down)[1]) + 1
85
    prefix = f"{version:0>3}."
86
    path.rename(path.parent / (prefix + path.name))
87
    print(f'Prefixed new migration script with "{prefix}".')
88
89
90
def run_migrations_offline():
91
    """Run migrations in 'offline' mode.
92
93
    This configures the context with just a URL
94
    and not an Engine, though an Engine is acceptable
95
    here as well.  By skipping the Engine creation
96
    we don't even need a DBAPI to be available.
97
98
    Calls to context.execute() here emit the given string to the
99
    script output.
100
101
    """
102
    url = config.get_main_option("sqlalchemy.url")
103
    context.configure(
104
        url=url,
105
        target_metadata=target_metadata,
106
        literal_binds=True,
107
        dialect_opts={"paramstyle": "named"},
108
    )
109
110
    with context.begin_transaction():
111
        context.run_migrations()
112
113
114
def run_migrations_online():
115
    """Run migrations in 'online' mode.
116
117
    In this scenario we need to create an Engine
118
    and associate a connection with the context.
119
120
    """
121
    connectable = engine()
122
123
    with connectable.connect() as connection:
124
        context.configure(
125
            connection=connection, target_metadata=target_metadata
126
        )
127
128
        with context.begin_transaction():
129
            context.run_migrations()
130
131
132
if context.is_offline_mode():
133
    run_migrations_offline()
134
else:
135
    run_migrations_online()
136