Coverage for /home/martinb/.local/share/virtualenvs/camcops/lib/python3.6/site-packages/cardinal_pythonlib/sqlalchemy/alembic_func.py : 28%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1#!/usr/bin/env python
2# cardinal_pythonlib/sqlalchemy/alembic_func.py
4"""
5===============================================================================
7 Original code copyright (C) 2009-2021 Rudolf Cardinal (rudolf@pobox.com).
9 This file is part of cardinal_pythonlib.
11 Licensed under the Apache License, Version 2.0 (the "License");
12 you may not use this file except in compliance with the License.
13 You may obtain a copy of the License at
15 https://www.apache.org/licenses/LICENSE-2.0
17 Unless required by applicable law or agreed to in writing, software
18 distributed under the License is distributed on an "AS IS" BASIS,
19 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
20 See the License for the specific language governing permissions and
21 limitations under the License.
23===============================================================================
25**Support functions for Alembic, the migration tool for SQLAlchemy.**
27"""
29import os
30import re
31import subprocess
32from typing import Tuple
34from alembic.config import Config
35from alembic.util.exc import CommandError
36from alembic.runtime.migration import MigrationContext
37from alembic.runtime.environment import EnvironmentContext
38from alembic.script import ScriptDirectory
39from sqlalchemy.engine import create_engine
41from cardinal_pythonlib.fileops import preserve_cwd
42from cardinal_pythonlib.logs import get_brace_style_log_with_null_handler
44log = get_brace_style_log_with_null_handler(__name__)
47# =============================================================================
48# Constants for Alembic
49# =============================================================================
50# https://alembic.readthedocs.org/en/latest/naming.html
51# http://docs.sqlalchemy.org/en/latest/core/constraints.html#configuring-constraint-naming-conventions # noqa
53ALEMBIC_NAMING_CONVENTION = {
54 "ix": 'ix_%(column_0_label)s',
55 "uq": "uq_%(table_name)s_%(column_0_name)s",
56 # "ck": "ck_%(table_name)s_%(constraint_name)s", # too long?
57 # ... https://groups.google.com/forum/#!topic/sqlalchemy/SIT4D8S9dUg
58 "ck": "ck_%(table_name)s_%(column_0_name)s",
59 "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
60 "pk": "pk_%(table_name)s"
61}
63DEFAULT_ALEMBIC_VERSION_TABLE = "alembic_version"
66# =============================================================================
67# Alembic revision/migration system
68# =============================================================================
69# https://stackoverflow.com/questions/24622170/using-alembic-api-from-inside-application-code # noqa
71def get_head_revision_from_alembic(
72 alembic_config_filename: str,
73 alembic_base_dir: str = None,
74 version_table: str = DEFAULT_ALEMBIC_VERSION_TABLE) -> str:
75 """
76 Ask Alembic what its head revision is (i.e. where the Python code would
77 like the database to be at).
79 Arguments:
80 alembic_config_filename: config filename
81 alembic_base_dir: directory to start in, so relative paths in the
82 config file work.
83 version_table: table name for Alembic versions
84 """
85 if alembic_base_dir is None:
86 alembic_base_dir = os.path.dirname(alembic_config_filename)
87 os.chdir(alembic_base_dir) # so the directory in the config file works
88 config = Config(alembic_config_filename)
89 script = ScriptDirectory.from_config(config)
90 with EnvironmentContext(config,
91 script,
92 version_table=version_table):
93 return script.get_current_head()
96def get_current_revision(
97 database_url: str,
98 version_table: str = DEFAULT_ALEMBIC_VERSION_TABLE) -> str:
99 """
100 Ask the database what its current revision is.
102 Arguments:
103 database_url: SQLAlchemy URL for the database
104 version_table: table name for Alembic versions
105 """
106 engine = create_engine(database_url)
107 conn = engine.connect()
108 opts = {'version_table': version_table}
109 mig_context = MigrationContext.configure(conn, opts=opts)
110 return mig_context.get_current_revision()
113def get_current_and_head_revision(
114 database_url: str,
115 alembic_config_filename: str,
116 alembic_base_dir: str = None,
117 version_table: str = DEFAULT_ALEMBIC_VERSION_TABLE) -> Tuple[str, str]:
118 """
119 Returns a tuple of ``(current_revision, head_revision)``; see
120 :func:`get_current_revision` and :func:`get_head_revision_from_alembic`.
122 Arguments:
123 database_url: SQLAlchemy URL for the database
124 alembic_config_filename: config filename
125 alembic_base_dir: directory to start in, so relative paths in the
126 config file work.
127 version_table: table name for Alembic versions
128 """
129 # Where we are
130 head_revision = get_head_revision_from_alembic(
131 alembic_config_filename=alembic_config_filename,
132 alembic_base_dir=alembic_base_dir,
133 version_table=version_table
134 )
135 log.info("Intended database version: {}", head_revision)
137 # Where we want to be
138 current_revision = get_current_revision(
139 database_url=database_url,
140 version_table=version_table
141 )
142 log.info("Current database version: {}", current_revision)
144 # Are we where we want to be?
145 return current_revision, head_revision
148@preserve_cwd
149def upgrade_database(
150 alembic_config_filename: str,
151 alembic_base_dir: str = None,
152 starting_revision: str = None,
153 destination_revision: str = "head",
154 version_table: str = DEFAULT_ALEMBIC_VERSION_TABLE,
155 as_sql: bool = False) -> None:
156 """
157 Use Alembic to upgrade our database.
159 See https://alembic.readthedocs.org/en/latest/api/runtime.html
160 but also, in particular, ``site-packages/alembic/command.py``
162 Arguments:
163 alembic_config_filename:
164 config filename
166 alembic_base_dir:
167 directory to start in, so relative paths in the config file work
169 starting_revision:
170 revision to start at (typically ``None`` to ask the database)
172 destination_revision:
173 revision to aim for (typically ``"head"`` to migrate to the latest
174 structure)
176 version_table: table name for Alembic versions
178 as_sql:
179 run in "offline" mode: print the migration SQL, rather than
180 modifying the database. See
181 https://alembic.zzzcomputing.com/en/latest/offline.html
183 """
185 if alembic_base_dir is None:
186 alembic_base_dir = os.path.dirname(alembic_config_filename)
187 os.chdir(alembic_base_dir) # so the directory in the config file works
188 config = Config(alembic_config_filename)
189 script = ScriptDirectory.from_config(config)
191 # noinspection PyUnusedLocal,PyProtectedMember
192 def upgrade(rev, context):
193 return script._upgrade_revs(destination_revision, rev)
195 log.info("Upgrading database to revision {!r} using Alembic",
196 destination_revision)
198 with EnvironmentContext(config,
199 script,
200 fn=upgrade,
201 as_sql=as_sql,
202 starting_rev=starting_revision,
203 destination_rev=destination_revision,
204 tag=None,
205 version_table=version_table):
206 script.run_env()
208 log.info("Database upgrade completed")
211@preserve_cwd
212def downgrade_database(
213 alembic_config_filename: str,
214 destination_revision: str,
215 alembic_base_dir: str = None,
216 starting_revision: str = None,
217 version_table: str = DEFAULT_ALEMBIC_VERSION_TABLE,
218 as_sql: bool = False) -> None:
219 """
220 Use Alembic to downgrade our database. USE WITH EXTREME CAUTION.
221 "revision" is the destination revision.
223 See https://alembic.readthedocs.org/en/latest/api/runtime.html
224 but also, in particular, ``site-packages/alembic/command.py``
226 Arguments:
227 alembic_config_filename:
228 config filename
230 alembic_base_dir:
231 directory to start in, so relative paths in the config file work
233 starting_revision:
234 revision to start at (typically ``None`` to ask the database)
236 destination_revision:
237 revision to aim for
239 version_table: table name for Alembic versions
241 as_sql:
242 run in "offline" mode: print the migration SQL, rather than
243 modifying the database. See
244 https://alembic.zzzcomputing.com/en/latest/offline.html
246 """
248 if alembic_base_dir is None:
249 alembic_base_dir = os.path.dirname(alembic_config_filename)
250 os.chdir(alembic_base_dir) # so the directory in the config file works
251 config = Config(alembic_config_filename)
252 script = ScriptDirectory.from_config(config)
254 # noinspection PyUnusedLocal,PyProtectedMember
255 def downgrade(rev, context):
256 return script._downgrade_revs(destination_revision, rev)
258 log.info("Downgrading database to revision {!r} using Alembic",
259 destination_revision)
261 with EnvironmentContext(config,
262 script,
263 fn=downgrade,
264 as_sql=as_sql,
265 starting_rev=starting_revision,
266 destination_rev=destination_revision,
267 tag=None,
268 version_table=version_table):
269 script.run_env()
271 log.info("Database downgrade completed")
274@preserve_cwd
275def create_database_migration_numbered_style(
276 alembic_ini_file: str,
277 alembic_versions_dir: str,
278 message: str,
279 n_sequence_chars: int = 4) -> None:
280 """
281 Create a new Alembic migration script.
283 Alembic compares the **state of the database** to the **state of the
284 metadata**, and generates a migration that brings the former up to the
285 latter. (It does **not** compare the most recent revision to the current
286 metadata, so make sure your database is up to date with the most recent
287 revision before running this!)
289 You **must check** that the autogenerated revisions are sensible.
291 How does it know where to look for the database?
293 1. This function changes into the directory of the Alembic ``.ini``
294 file and calls the external program
296 .. code-block:: bash
298 alembic -c ALEMBIC_INI_FILE revision --autogenerate -m MESSAGE --rev-id REVISION_ID
300 2. The Alembic ``.ini`` file points (via the ``script_location``
301 variable) to a directory containing your ``env.py``. Alembic loads
302 this script.
304 3. That script typically works out the database URL and calls further
305 into the Alembic code.
307 See https://alembic.zzzcomputing.com/en/latest/autogenerate.html.
309 Regarding filenames: the default ``n_sequence_chars`` of 4 is like Django
310 and gives files with names like
312 .. code-block:: none
314 0001_x.py, 0002_y.py, ...
316 NOTE THAT TO USE A NON-STANDARD ALEMBIC VERSION TABLE, YOU MUST SPECIFY
317 THAT IN YOUR ``env.py`` (see e.g. CamCOPS).
319 Args:
320 alembic_ini_file: filename of Alembic ``alembic.ini`` file
321 alembic_versions_dir: directory in which you keep your Python scripts,
322 one per Alembic revision
323 message: message to be associated with this revision
324 n_sequence_chars: number of numerical sequence characters to use in the
325 filename/revision (see above).
326 """ # noqa
327 file_regex = r"\d{" + str(n_sequence_chars) + r"}_\S*\.py$"
329 _, _, existing_version_filenames = next(os.walk(alembic_versions_dir),
330 (None, None, []))
331 existing_version_filenames = [
332 x for x in existing_version_filenames if re.match(file_regex, x)]
333 log.debug("Existing Alembic version script filenames: {!r}",
334 existing_version_filenames)
335 current_seq_strs = [x[:n_sequence_chars]
336 for x in existing_version_filenames]
337 current_seq_strs.sort()
338 if not current_seq_strs:
339 current_seq_str = None
340 new_seq_no = 1
341 else:
342 current_seq_str = current_seq_strs[-1]
343 new_seq_no = max(int(x) for x in current_seq_strs) + 1
344 new_seq_str = str(new_seq_no).zfill(n_sequence_chars)
346 log.info(
347 """
348Generating new revision with Alembic...
349 Last revision was: {}
350 New revision will be: {}
351 [If it fails with "Can't locate revision identified by...", you might need
352 to DROP the Alembic version table (by default named 'alembic_version', but
353 you may have elected to change that in your env.py.]
354 """,
355 current_seq_str,
356 new_seq_str
357 )
359 alembic_ini_dir = os.path.dirname(alembic_ini_file)
360 os.chdir(alembic_ini_dir)
361 cmdargs = ['alembic',
362 '-c', alembic_ini_file,
363 'revision',
364 '--autogenerate',
365 '-m', message,
366 '--rev-id', new_seq_str]
367 log.info("From directory {!r}, calling: {!r}", alembic_ini_dir, cmdargs)
368 subprocess.call(cmdargs)
371def stamp_allowing_unusual_version_table(
372 config: Config,
373 revision: str,
374 sql: bool = False,
375 tag: str = None,
376 version_table: str = DEFAULT_ALEMBIC_VERSION_TABLE) -> None:
377 """
378 Stamps the Alembic version table with the given revision; don't run any
379 migrations.
381 This function is a clone of ``alembic.command.stamp()``, but allowing
382 ``version_table`` to change. See
383 https://alembic.zzzcomputing.com/en/latest/api/commands.html#alembic.command.stamp
384 """ # noqa
386 script = ScriptDirectory.from_config(config)
388 starting_rev = None
389 if ":" in revision:
390 if not sql:
391 raise CommandError("Range revision not allowed")
392 starting_rev, revision = revision.split(':', 2)
394 # noinspection PyUnusedLocal
395 def do_stamp(rev: str, context):
396 # noinspection PyProtectedMember
397 return script._stamp_revs(revision, rev)
399 with EnvironmentContext(config,
400 script,
401 fn=do_stamp,
402 as_sql=sql,
403 destination_rev=revision,
404 starting_rev=starting_rev,
405 tag=tag,
406 version_table=version_table):
407 script.run_env()