Hide keyboard shortcuts

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 

3 

4""" 

5=============================================================================== 

6 

7 Original code copyright (C) 2009-2021 Rudolf Cardinal (rudolf@pobox.com). 

8 

9 This file is part of cardinal_pythonlib. 

10 

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 

14 

15 https://www.apache.org/licenses/LICENSE-2.0 

16 

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. 

22 

23=============================================================================== 

24 

25**Support functions for Alembic, the migration tool for SQLAlchemy.** 

26 

27""" 

28 

29import os 

30import re 

31import subprocess 

32from typing import Tuple 

33 

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 

40 

41from cardinal_pythonlib.fileops import preserve_cwd 

42from cardinal_pythonlib.logs import get_brace_style_log_with_null_handler 

43 

44log = get_brace_style_log_with_null_handler(__name__) 

45 

46 

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 

52 

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} 

62 

63DEFAULT_ALEMBIC_VERSION_TABLE = "alembic_version" 

64 

65 

66# ============================================================================= 

67# Alembic revision/migration system 

68# ============================================================================= 

69# https://stackoverflow.com/questions/24622170/using-alembic-api-from-inside-application-code # noqa 

70 

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). 

78 

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() 

94 

95 

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. 

101 

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() 

111 

112 

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`. 

121 

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) 

136 

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) 

143 

144 # Are we where we want to be? 

145 return current_revision, head_revision 

146 

147 

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. 

158 

159 See https://alembic.readthedocs.org/en/latest/api/runtime.html 

160 but also, in particular, ``site-packages/alembic/command.py`` 

161 

162 Arguments: 

163 alembic_config_filename: 

164 config filename 

165 

166 alembic_base_dir: 

167 directory to start in, so relative paths in the config file work 

168 

169 starting_revision: 

170 revision to start at (typically ``None`` to ask the database) 

171 

172 destination_revision: 

173 revision to aim for (typically ``"head"`` to migrate to the latest 

174 structure) 

175 

176 version_table: table name for Alembic versions 

177 

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 

182 

183 """ 

184 

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) 

190 

191 # noinspection PyUnusedLocal,PyProtectedMember 

192 def upgrade(rev, context): 

193 return script._upgrade_revs(destination_revision, rev) 

194 

195 log.info("Upgrading database to revision {!r} using Alembic", 

196 destination_revision) 

197 

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() 

207 

208 log.info("Database upgrade completed") 

209 

210 

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. 

222 

223 See https://alembic.readthedocs.org/en/latest/api/runtime.html 

224 but also, in particular, ``site-packages/alembic/command.py`` 

225 

226 Arguments: 

227 alembic_config_filename: 

228 config filename 

229 

230 alembic_base_dir: 

231 directory to start in, so relative paths in the config file work 

232 

233 starting_revision: 

234 revision to start at (typically ``None`` to ask the database) 

235 

236 destination_revision: 

237 revision to aim for 

238 

239 version_table: table name for Alembic versions 

240 

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 

245 

246 """ 

247 

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) 

253 

254 # noinspection PyUnusedLocal,PyProtectedMember 

255 def downgrade(rev, context): 

256 return script._downgrade_revs(destination_revision, rev) 

257 

258 log.info("Downgrading database to revision {!r} using Alembic", 

259 destination_revision) 

260 

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() 

270 

271 log.info("Database downgrade completed") 

272 

273 

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. 

282 

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!) 

288 

289 You **must check** that the autogenerated revisions are sensible. 

290 

291 How does it know where to look for the database? 

292 

293 1. This function changes into the directory of the Alembic ``.ini`` 

294 file and calls the external program 

295 

296 .. code-block:: bash 

297 

298 alembic -c ALEMBIC_INI_FILE revision --autogenerate -m MESSAGE --rev-id REVISION_ID 

299 

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. 

303 

304 3. That script typically works out the database URL and calls further 

305 into the Alembic code. 

306 

307 See https://alembic.zzzcomputing.com/en/latest/autogenerate.html. 

308 

309 Regarding filenames: the default ``n_sequence_chars`` of 4 is like Django 

310 and gives files with names like 

311 

312 .. code-block:: none 

313 

314 0001_x.py, 0002_y.py, ... 

315 

316 NOTE THAT TO USE A NON-STANDARD ALEMBIC VERSION TABLE, YOU MUST SPECIFY 

317 THAT IN YOUR ``env.py`` (see e.g. CamCOPS). 

318 

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$" 

328 

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) 

345 

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 ) 

358 

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) 

369 

370 

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. 

380 

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 

385 

386 script = ScriptDirectory.from_config(config) 

387 

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) 

393 

394 # noinspection PyUnusedLocal 

395 def do_stamp(rev: str, context): 

396 # noinspection PyProtectedMember 

397 return script._stamp_revs(revision, rev) 

398 

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()