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

1from contextlib import contextmanager 

2import logging 

3import sys 

4 

5from sqlalchemy import Column 

6from sqlalchemy import literal_column 

7from sqlalchemy import MetaData 

8from sqlalchemy import PrimaryKeyConstraint 

9from sqlalchemy import String 

10from sqlalchemy import Table 

11from sqlalchemy.engine import Connection 

12from sqlalchemy.engine import url as sqla_url 

13from sqlalchemy.engine.strategies import MockEngineStrategy 

14 

15from .. import ddl 

16from .. import util 

17from ..util import sqla_compat 

18from ..util.compat import callable 

19from ..util.compat import EncodedIO 

20 

21log = logging.getLogger(__name__) 

22 

23 

24class _ProxyTransaction(object): 

25 def __init__(self, migration_context): 

26 self.migration_context = migration_context 

27 

28 @property 

29 def _proxied_transaction(self): 

30 return self.migration_context._transaction 

31 

32 def rollback(self): 

33 self._proxied_transaction.rollback() 

34 

35 def commit(self): 

36 self._proxied_transaction.commit() 

37 

38 def __enter__(self): 

39 return self 

40 

41 def __exit__(self, type_, value, traceback): 

42 self._proxied_transaction.__exit__(type_, value, traceback) 

43 

44 

45class MigrationContext(object): 

46 

47 """Represent the database state made available to a migration 

48 script. 

49 

50 :class:`.MigrationContext` is the front end to an actual 

51 database connection, or alternatively a string output 

52 stream given a particular database dialect, 

53 from an Alembic perspective. 

54 

55 When inside the ``env.py`` script, the :class:`.MigrationContext` 

56 is available via the 

57 :meth:`.EnvironmentContext.get_context` method, 

58 which is available at ``alembic.context``:: 

59 

60 # from within env.py script 

61 from alembic import context 

62 migration_context = context.get_context() 

63 

64 For usage outside of an ``env.py`` script, such as for 

65 utility routines that want to check the current version 

66 in the database, the :meth:`.MigrationContext.configure` 

67 method to create new :class:`.MigrationContext` objects. 

68 For example, to get at the current revision in the 

69 database using :meth:`.MigrationContext.get_current_revision`:: 

70 

71 # in any application, outside of an env.py script 

72 from alembic.migration import MigrationContext 

73 from sqlalchemy import create_engine 

74 

75 engine = create_engine("postgresql://mydatabase") 

76 conn = engine.connect() 

77 

78 context = MigrationContext.configure(conn) 

79 current_rev = context.get_current_revision() 

80 

81 The above context can also be used to produce 

82 Alembic migration operations with an :class:`.Operations` 

83 instance:: 

84 

85 # in any application, outside of the normal Alembic environment 

86 from alembic.operations import Operations 

87 op = Operations(context) 

88 op.alter_column("mytable", "somecolumn", nullable=True) 

89 

90 """ 

91 

92 def __init__(self, dialect, connection, opts, environment_context=None): 

93 self.environment_context = environment_context 

94 self.opts = opts 

95 self.dialect = dialect 

96 self.script = opts.get("script") 

97 as_sql = opts.get("as_sql", False) 

98 transactional_ddl = opts.get("transactional_ddl") 

99 self._transaction_per_migration = opts.get( 

100 "transaction_per_migration", False 

101 ) 

102 self.on_version_apply_callbacks = opts.get("on_version_apply", ()) 

103 self._transaction = None 

104 

105 if as_sql: 

106 self.connection = self._stdout_connection(connection) 

107 assert self.connection is not None 

108 else: 

109 self.connection = connection 

110 self._migrations_fn = opts.get("fn") 

111 self.as_sql = as_sql 

112 

113 self.purge = opts.get("purge", False) 

114 

115 if "output_encoding" in opts: 

116 self.output_buffer = EncodedIO( 

117 opts.get("output_buffer") or sys.stdout, 

118 opts["output_encoding"], 

119 ) 

120 else: 

121 self.output_buffer = opts.get("output_buffer", sys.stdout) 

122 

123 self._user_compare_type = opts.get("compare_type", False) 

124 self._user_compare_server_default = opts.get( 

125 "compare_server_default", False 

126 ) 

127 self.version_table = version_table = opts.get( 

128 "version_table", "alembic_version" 

129 ) 

130 self.version_table_schema = version_table_schema = opts.get( 

131 "version_table_schema", None 

132 ) 

133 self._version = Table( 

134 version_table, 

135 MetaData(), 

136 Column("version_num", String(32), nullable=False), 

137 schema=version_table_schema, 

138 ) 

139 if opts.get("version_table_pk", True): 

140 self._version.append_constraint( 

141 PrimaryKeyConstraint( 

142 "version_num", name="%s_pkc" % version_table 

143 ) 

144 ) 

145 

146 self._start_from_rev = opts.get("starting_rev") 

147 self.impl = ddl.DefaultImpl.get_by_dialect(dialect)( 

148 dialect, 

149 self.connection, 

150 self.as_sql, 

151 transactional_ddl, 

152 self.output_buffer, 

153 opts, 

154 ) 

155 log.info("Context impl %s.", self.impl.__class__.__name__) 

156 if self.as_sql: 

157 log.info("Generating static SQL") 

158 log.info( 

159 "Will assume %s DDL.", 

160 "transactional" 

161 if self.impl.transactional_ddl 

162 else "non-transactional", 

163 ) 

164 

165 @classmethod 

166 def configure( 

167 cls, 

168 connection=None, 

169 url=None, 

170 dialect_name=None, 

171 dialect=None, 

172 environment_context=None, 

173 dialect_opts=None, 

174 opts=None, 

175 ): 

176 """Create a new :class:`.MigrationContext`. 

177 

178 This is a factory method usually called 

179 by :meth:`.EnvironmentContext.configure`. 

180 

181 :param connection: a :class:`~sqlalchemy.engine.Connection` 

182 to use for SQL execution in "online" mode. When present, 

183 is also used to determine the type of dialect in use. 

184 :param url: a string database url, or a 

185 :class:`sqlalchemy.engine.url.URL` object. 

186 The type of dialect to be used will be derived from this if 

187 ``connection`` is not passed. 

188 :param dialect_name: string name of a dialect, such as 

189 "postgresql", "mssql", etc. The type of dialect to be used will be 

190 derived from this if ``connection`` and ``url`` are not passed. 

191 :param opts: dictionary of options. Most other options 

192 accepted by :meth:`.EnvironmentContext.configure` are passed via 

193 this dictionary. 

194 

195 """ 

196 if opts is None: 

197 opts = {} 

198 if dialect_opts is None: 

199 dialect_opts = {} 

200 

201 if connection: 

202 if not isinstance(connection, Connection): 

203 util.warn( 

204 "'connection' argument to configure() is expected " 

205 "to be a sqlalchemy.engine.Connection instance, " 

206 "got %r" % connection, 

207 stacklevel=3, 

208 ) 

209 

210 dialect = connection.dialect 

211 elif url: 

212 url = sqla_url.make_url(url) 

213 dialect = url.get_dialect()(**dialect_opts) 

214 elif dialect_name: 

215 url = sqla_url.make_url("%s://" % dialect_name) 

216 dialect = url.get_dialect()(**dialect_opts) 

217 elif not dialect: 

218 raise Exception("Connection, url, or dialect_name is required.") 

219 

220 return MigrationContext(dialect, connection, opts, environment_context) 

221 

222 @contextmanager 

223 def autocommit_block(self): 

224 """Enter an "autocommit" block, for databases that support AUTOCOMMIT 

225 isolation levels. 

226 

227 This special directive is intended to support the occasional database 

228 DDL or system operation that specifically has to be run outside of 

229 any kind of transaction block. The PostgreSQL database platform 

230 is the most common target for this style of operation, as many 

231 of its DDL operations must be run outside of transaction blocks, even 

232 though the database overall supports transactional DDL. 

233 

234 The method is used as a context manager within a migration script, by 

235 calling on :meth:`.Operations.get_context` to retrieve the 

236 :class:`.MigrationContext`, then invoking 

237 :meth:`.MigrationContext.autocommit_block` using the ``with:`` 

238 statement:: 

239 

240 def upgrade(): 

241 with op.get_context().autocommit_block(): 

242 op.execute("ALTER TYPE mood ADD VALUE 'soso'") 

243 

244 Above, a PostgreSQL "ALTER TYPE..ADD VALUE" directive is emitted, 

245 which must be run outside of a transaction block at the database level. 

246 The :meth:`.MigrationContext.autocommit_block` method makes use of the 

247 SQLAlchemy ``AUTOCOMMIT`` isolation level setting, which against the 

248 psycogp2 DBAPI corresponds to the ``connection.autocommit`` setting, 

249 to ensure that the database driver is not inside of a DBAPI level 

250 transaction block. 

251 

252 .. warning:: 

253 

254 As is necessary, **the database transaction preceding the block is 

255 unconditionally committed**. This means that the run of migrations 

256 preceding the operation will be committed, before the overall 

257 migration operation is complete. 

258 

259 It is recommended that when an application includes migrations with 

260 "autocommit" blocks, that 

261 :paramref:`.EnvironmentContext.transaction_per_migration` be used 

262 so that the calling environment is tuned to expect short per-file 

263 migrations whether or not one of them has an autocommit block. 

264 

265 

266 .. versionadded:: 1.2.0 

267 

268 """ 

269 _in_connection_transaction = self._in_connection_transaction() 

270 

271 if self.impl.transactional_ddl: 

272 if self.as_sql: 

273 self.impl.emit_commit() 

274 

275 elif _in_connection_transaction: 

276 assert self._transaction is not None 

277 

278 self._transaction.commit() 

279 self._transaction = None 

280 

281 if not self.as_sql: 

282 current_level = self.connection.get_isolation_level() 

283 self.connection.execution_options(isolation_level="AUTOCOMMIT") 

284 try: 

285 yield 

286 finally: 

287 if not self.as_sql: 

288 self.connection.execution_options( 

289 isolation_level=current_level 

290 ) 

291 

292 if self.impl.transactional_ddl: 

293 if self.as_sql: 

294 self.impl.emit_begin() 

295 

296 elif _in_connection_transaction: 

297 self._transaction = self.bind.begin() 

298 

299 def begin_transaction(self, _per_migration=False): 

300 """Begin a logical transaction for migration operations. 

301 

302 This method is used within an ``env.py`` script to demarcate where 

303 the outer "transaction" for a series of migrations begins. Example:: 

304 

305 def run_migrations_online(): 

306 connectable = create_engine(...) 

307 

308 with connectable.connect() as connection: 

309 context.configure( 

310 connection=connection, target_metadata=target_metadata 

311 ) 

312 

313 with context.begin_transaction(): 

314 context.run_migrations() 

315 

316 Above, :meth:`.MigrationContext.begin_transaction` is used to demarcate 

317 where the outer logical transaction occurs around the 

318 :meth:`.MigrationContext.run_migrations` operation. 

319 

320 A "Logical" transaction means that the operation may or may not 

321 correspond to a real database transaction. If the target database 

322 supports transactional DDL (or 

323 :paramref:`.EnvironmentContext.configure.transactional_ddl` is true), 

324 the :paramref:`.EnvironmentContext.configure.transaction_per_migration` 

325 flag is not set, and the migration is against a real database 

326 connection (as opposed to using "offline" ``--sql`` mode), a real 

327 transaction will be started. If ``--sql`` mode is in effect, the 

328 operation would instead correspond to a string such as "BEGIN" being 

329 emitted to the string output. 

330 

331 The returned object is a Python context manager that should only be 

332 used in the context of a ``with:`` statement as indicated above. 

333 The object has no other guaranteed API features present. 

334 

335 .. seealso:: 

336 

337 :meth:`.MigrationContext.autocommit_block` 

338 

339 """ 

340 transaction_now = _per_migration == self._transaction_per_migration 

341 

342 if not transaction_now: 

343 

344 @contextmanager 

345 def do_nothing(): 

346 yield 

347 

348 return do_nothing() 

349 

350 elif not self.impl.transactional_ddl: 

351 

352 @contextmanager 

353 def do_nothing(): 

354 yield 

355 

356 return do_nothing() 

357 elif self.as_sql: 

358 

359 @contextmanager 

360 def begin_commit(): 

361 self.impl.emit_begin() 

362 yield 

363 self.impl.emit_commit() 

364 

365 return begin_commit() 

366 else: 

367 self._transaction = self.bind.begin() 

368 return _ProxyTransaction(self) 

369 

370 def get_current_revision(self): 

371 """Return the current revision, usually that which is present 

372 in the ``alembic_version`` table in the database. 

373 

374 This method intends to be used only for a migration stream that 

375 does not contain unmerged branches in the target database; 

376 if there are multiple branches present, an exception is raised. 

377 The :meth:`.MigrationContext.get_current_heads` should be preferred 

378 over this method going forward in order to be compatible with 

379 branch migration support. 

380 

381 If this :class:`.MigrationContext` was configured in "offline" 

382 mode, that is with ``as_sql=True``, the ``starting_rev`` 

383 parameter is returned instead, if any. 

384 

385 """ 

386 heads = self.get_current_heads() 

387 if len(heads) == 0: 

388 return None 

389 elif len(heads) > 1: 

390 raise util.CommandError( 

391 "Version table '%s' has more than one head present; " 

392 "please use get_current_heads()" % self.version_table 

393 ) 

394 else: 

395 return heads[0] 

396 

397 def get_current_heads(self): 

398 """Return a tuple of the current 'head versions' that are represented 

399 in the target database. 

400 

401 For a migration stream without branches, this will be a single 

402 value, synonymous with that of 

403 :meth:`.MigrationContext.get_current_revision`. However when multiple 

404 unmerged branches exist within the target database, the returned tuple 

405 will contain a value for each head. 

406 

407 If this :class:`.MigrationContext` was configured in "offline" 

408 mode, that is with ``as_sql=True``, the ``starting_rev`` 

409 parameter is returned in a one-length tuple. 

410 

411 If no version table is present, or if there are no revisions 

412 present, an empty tuple is returned. 

413 

414 .. versionadded:: 0.7.0 

415 

416 """ 

417 if self.as_sql: 

418 start_from_rev = self._start_from_rev 

419 if start_from_rev == "base": 

420 start_from_rev = None 

421 elif start_from_rev is not None and self.script: 

422 

423 start_from_rev = [ 

424 self.script.get_revision(sfr).revision 

425 for sfr in util.to_list(start_from_rev) 

426 if sfr not in (None, "base") 

427 ] 

428 return util.to_tuple(start_from_rev, default=()) 

429 else: 

430 if self._start_from_rev: 

431 raise util.CommandError( 

432 "Can't specify current_rev to context " 

433 "when using a database connection" 

434 ) 

435 if not self._has_version_table(): 

436 return () 

437 return tuple( 

438 row[0] for row in self.connection.execute(self._version.select()) 

439 ) 

440 

441 def _ensure_version_table(self, purge=False): 

442 self._version.create(self.connection, checkfirst=True) 

443 if purge: 

444 self.connection.execute(self._version.delete()) 

445 

446 def _has_version_table(self): 

447 return sqla_compat._connectable_has_table( 

448 self.connection, self.version_table, self.version_table_schema 

449 ) 

450 

451 def stamp(self, script_directory, revision): 

452 """Stamp the version table with a specific revision. 

453 

454 This method calculates those branches to which the given revision 

455 can apply, and updates those branches as though they were migrated 

456 towards that revision (either up or down). If no current branches 

457 include the revision, it is added as a new branch head. 

458 

459 .. versionadded:: 0.7.0 

460 

461 """ 

462 heads = self.get_current_heads() 

463 if not self.as_sql and not heads: 

464 self._ensure_version_table() 

465 head_maintainer = HeadMaintainer(self, heads) 

466 for step in script_directory._stamp_revs(revision, heads): 

467 head_maintainer.update_to_step(step) 

468 

469 def run_migrations(self, **kw): 

470 r"""Run the migration scripts established for this 

471 :class:`.MigrationContext`, if any. 

472 

473 The commands in :mod:`alembic.command` will set up a function 

474 that is ultimately passed to the :class:`.MigrationContext` 

475 as the ``fn`` argument. This function represents the "work" 

476 that will be done when :meth:`.MigrationContext.run_migrations` 

477 is called, typically from within the ``env.py`` script of the 

478 migration environment. The "work function" then provides an iterable 

479 of version callables and other version information which 

480 in the case of the ``upgrade`` or ``downgrade`` commands are the 

481 list of version scripts to invoke. Other commands yield nothing, 

482 in the case that a command wants to run some other operation 

483 against the database such as the ``current`` or ``stamp`` commands. 

484 

485 :param \**kw: keyword arguments here will be passed to each 

486 migration callable, that is the ``upgrade()`` or ``downgrade()`` 

487 method within revision scripts. 

488 

489 """ 

490 self.impl.start_migrations() 

491 

492 if self.purge: 

493 if self.as_sql: 

494 raise util.CommandError("Can't use --purge with --sql mode") 

495 self._ensure_version_table(purge=True) 

496 heads = () 

497 else: 

498 heads = self.get_current_heads() 

499 

500 if not self.as_sql and not heads: 

501 self._ensure_version_table() 

502 

503 head_maintainer = HeadMaintainer(self, heads) 

504 

505 starting_in_transaction = ( 

506 not self.as_sql and self._in_connection_transaction() 

507 ) 

508 

509 for step in self._migrations_fn(heads, self): 

510 with self.begin_transaction(_per_migration=True): 

511 if self.as_sql and not head_maintainer.heads: 

512 # for offline mode, include a CREATE TABLE from 

513 # the base 

514 self._version.create(self.connection) 

515 log.info("Running %s", step) 

516 if self.as_sql: 

517 self.impl.static_output( 

518 "-- Running %s" % (step.short_log,) 

519 ) 

520 step.migration_fn(**kw) 

521 

522 # previously, we wouldn't stamp per migration 

523 # if we were in a transaction, however given the more 

524 # complex model that involves any number of inserts 

525 # and row-targeted updates and deletes, it's simpler for now 

526 # just to run the operations on every version 

527 head_maintainer.update_to_step(step) 

528 for callback in self.on_version_apply_callbacks: 

529 callback( 

530 ctx=self, 

531 step=step.info, 

532 heads=set(head_maintainer.heads), 

533 run_args=kw, 

534 ) 

535 

536 if ( 

537 not starting_in_transaction 

538 and not self.as_sql 

539 and not self.impl.transactional_ddl 

540 and self._in_connection_transaction() 

541 ): 

542 raise util.CommandError( 

543 'Migration "%s" has left an uncommitted ' 

544 "transaction opened; transactional_ddl is False so " 

545 "Alembic is not committing transactions" % step 

546 ) 

547 

548 if self.as_sql and not head_maintainer.heads: 

549 self._version.drop(self.connection) 

550 

551 def _in_connection_transaction(self): 

552 try: 

553 meth = self.connection.in_transaction 

554 except AttributeError: 

555 return False 

556 else: 

557 return meth() 

558 

559 def execute(self, sql, execution_options=None): 

560 """Execute a SQL construct or string statement. 

561 

562 The underlying execution mechanics are used, that is 

563 if this is "offline mode" the SQL is written to the 

564 output buffer, otherwise the SQL is emitted on 

565 the current SQLAlchemy connection. 

566 

567 """ 

568 self.impl._exec(sql, execution_options) 

569 

570 def _stdout_connection(self, connection): 

571 def dump(construct, *multiparams, **params): 

572 self.impl._exec(construct) 

573 

574 return MockEngineStrategy.MockConnection(self.dialect, dump) 

575 

576 @property 

577 def bind(self): 

578 """Return the current "bind". 

579 

580 In online mode, this is an instance of 

581 :class:`sqlalchemy.engine.Connection`, and is suitable 

582 for ad-hoc execution of any kind of usage described 

583 in :ref:`sqlexpression_toplevel` as well as 

584 for usage with the :meth:`sqlalchemy.schema.Table.create` 

585 and :meth:`sqlalchemy.schema.MetaData.create_all` methods 

586 of :class:`~sqlalchemy.schema.Table`, 

587 :class:`~sqlalchemy.schema.MetaData`. 

588 

589 Note that when "standard output" mode is enabled, 

590 this bind will be a "mock" connection handler that cannot 

591 return results and is only appropriate for a very limited 

592 subset of commands. 

593 

594 """ 

595 return self.connection 

596 

597 @property 

598 def config(self): 

599 """Return the :class:`.Config` used by the current environment, if any. 

600 

601 .. versionadded:: 0.6.6 

602 

603 """ 

604 if self.environment_context: 

605 return self.environment_context.config 

606 else: 

607 return None 

608 

609 def _compare_type(self, inspector_column, metadata_column): 

610 if self._user_compare_type is False: 

611 return False 

612 

613 if callable(self._user_compare_type): 

614 user_value = self._user_compare_type( 

615 self, 

616 inspector_column, 

617 metadata_column, 

618 inspector_column.type, 

619 metadata_column.type, 

620 ) 

621 if user_value is not None: 

622 return user_value 

623 

624 return self.impl.compare_type(inspector_column, metadata_column) 

625 

626 def _compare_server_default( 

627 self, 

628 inspector_column, 

629 metadata_column, 

630 rendered_metadata_default, 

631 rendered_column_default, 

632 ): 

633 

634 if self._user_compare_server_default is False: 

635 return False 

636 

637 if callable(self._user_compare_server_default): 

638 user_value = self._user_compare_server_default( 

639 self, 

640 inspector_column, 

641 metadata_column, 

642 rendered_column_default, 

643 metadata_column.server_default, 

644 rendered_metadata_default, 

645 ) 

646 if user_value is not None: 

647 return user_value 

648 

649 return self.impl.compare_server_default( 

650 inspector_column, 

651 metadata_column, 

652 rendered_metadata_default, 

653 rendered_column_default, 

654 ) 

655 

656 

657class HeadMaintainer(object): 

658 def __init__(self, context, heads): 

659 self.context = context 

660 self.heads = set(heads) 

661 

662 def _insert_version(self, version): 

663 assert version not in self.heads 

664 self.heads.add(version) 

665 

666 self.context.impl._exec( 

667 self.context._version.insert().values( 

668 version_num=literal_column("'%s'" % version) 

669 ) 

670 ) 

671 

672 def _delete_version(self, version): 

673 self.heads.remove(version) 

674 

675 ret = self.context.impl._exec( 

676 self.context._version.delete().where( 

677 self.context._version.c.version_num 

678 == literal_column("'%s'" % version) 

679 ) 

680 ) 

681 if ( 

682 not self.context.as_sql 

683 and self.context.dialect.supports_sane_rowcount 

684 and ret.rowcount != 1 

685 ): 

686 raise util.CommandError( 

687 "Online migration expected to match one " 

688 "row when deleting '%s' in '%s'; " 

689 "%d found" 

690 % (version, self.context.version_table, ret.rowcount) 

691 ) 

692 

693 def _update_version(self, from_, to_): 

694 assert to_ not in self.heads 

695 self.heads.remove(from_) 

696 self.heads.add(to_) 

697 

698 ret = self.context.impl._exec( 

699 self.context._version.update() 

700 .values(version_num=literal_column("'%s'" % to_)) 

701 .where( 

702 self.context._version.c.version_num 

703 == literal_column("'%s'" % from_) 

704 ) 

705 ) 

706 if ( 

707 not self.context.as_sql 

708 and self.context.dialect.supports_sane_rowcount 

709 and ret.rowcount != 1 

710 ): 

711 raise util.CommandError( 

712 "Online migration expected to match one " 

713 "row when updating '%s' to '%s' in '%s'; " 

714 "%d found" 

715 % (from_, to_, self.context.version_table, ret.rowcount) 

716 ) 

717 

718 def update_to_step(self, step): 

719 if step.should_delete_branch(self.heads): 

720 vers = step.delete_version_num 

721 log.debug("branch delete %s", vers) 

722 self._delete_version(vers) 

723 elif step.should_create_branch(self.heads): 

724 vers = step.insert_version_num 

725 log.debug("new branch insert %s", vers) 

726 self._insert_version(vers) 

727 elif step.should_merge_branches(self.heads): 

728 # delete revs, update from rev, update to rev 

729 ( 

730 delete_revs, 

731 update_from_rev, 

732 update_to_rev, 

733 ) = step.merge_branch_idents(self.heads) 

734 log.debug( 

735 "merge, delete %s, update %s to %s", 

736 delete_revs, 

737 update_from_rev, 

738 update_to_rev, 

739 ) 

740 for delrev in delete_revs: 

741 self._delete_version(delrev) 

742 self._update_version(update_from_rev, update_to_rev) 

743 elif step.should_unmerge_branches(self.heads): 

744 ( 

745 update_from_rev, 

746 update_to_rev, 

747 insert_revs, 

748 ) = step.unmerge_branch_idents(self.heads) 

749 log.debug( 

750 "unmerge, insert %s, update %s to %s", 

751 insert_revs, 

752 update_from_rev, 

753 update_to_rev, 

754 ) 

755 for insrev in insert_revs: 

756 self._insert_version(insrev) 

757 self._update_version(update_from_rev, update_to_rev) 

758 else: 

759 from_, to_ = step.update_version_num(self.heads) 

760 log.debug("update %s to %s", from_, to_) 

761 self._update_version(from_, to_) 

762 

763 

764class MigrationInfo(object): 

765 """Exposes information about a migration step to a callback listener. 

766 

767 The :class:`.MigrationInfo` object is available exclusively for the 

768 benefit of the :paramref:`.EnvironmentContext.on_version_apply` 

769 callback hook. 

770 

771 .. versionadded:: 0.9.3 

772 

773 """ 

774 

775 is_upgrade = None 

776 """True/False: indicates whether this operation ascends or descends the 

777 version tree.""" 

778 

779 is_stamp = None 

780 """True/False: indicates whether this operation is a stamp (i.e. whether 

781 it results in any actual database operations).""" 

782 

783 up_revision_id = None 

784 """Version string corresponding to :attr:`.Revision.revision`. 

785 

786 In the case of a stamp operation, it is advised to use the 

787 :attr:`.MigrationInfo.up_revision_ids` tuple as a stamp operation can 

788 make a single movement from one or more branches down to a single 

789 branchpoint, in which case there will be multiple "up" revisions. 

790 

791 .. seealso:: 

792 

793 :attr:`.MigrationInfo.up_revision_ids` 

794 

795 """ 

796 

797 up_revision_ids = None 

798 """Tuple of version strings corresponding to :attr:`.Revision.revision`. 

799 

800 In the majority of cases, this tuple will be a single value, synonomous 

801 with the scalar value of :attr:`.MigrationInfo.up_revision_id`. 

802 It can be multiple revision identifiers only in the case of an 

803 ``alembic stamp`` operation which is moving downwards from multiple 

804 branches down to their common branch point. 

805 

806 .. versionadded:: 0.9.4 

807 

808 """ 

809 

810 down_revision_ids = None 

811 """Tuple of strings representing the base revisions of this migration step. 

812 

813 If empty, this represents a root revision; otherwise, the first item 

814 corresponds to :attr:`.Revision.down_revision`, and the rest are inferred 

815 from dependencies. 

816 """ 

817 

818 revision_map = None 

819 """The revision map inside of which this operation occurs.""" 

820 

821 def __init__( 

822 self, revision_map, is_upgrade, is_stamp, up_revisions, down_revisions 

823 ): 

824 self.revision_map = revision_map 

825 self.is_upgrade = is_upgrade 

826 self.is_stamp = is_stamp 

827 self.up_revision_ids = util.to_tuple(up_revisions, default=()) 

828 if self.up_revision_ids: 

829 self.up_revision_id = self.up_revision_ids[0] 

830 else: 

831 # this should never be the case with 

832 # "upgrade", "downgrade", or "stamp" as we are always 

833 # measuring movement in terms of at least one upgrade version 

834 self.up_revision_id = None 

835 self.down_revision_ids = util.to_tuple(down_revisions, default=()) 

836 

837 @property 

838 def is_migration(self): 

839 """True/False: indicates whether this operation is a migration. 

840 

841 At present this is true if and only the migration is not a stamp. 

842 If other operation types are added in the future, both this attribute 

843 and :attr:`~.MigrationInfo.is_stamp` will be false. 

844 """ 

845 return not self.is_stamp 

846 

847 @property 

848 def source_revision_ids(self): 

849 """Active revisions before this migration step is applied.""" 

850 return ( 

851 self.down_revision_ids if self.is_upgrade else self.up_revision_ids 

852 ) 

853 

854 @property 

855 def destination_revision_ids(self): 

856 """Active revisions after this migration step is applied.""" 

857 return ( 

858 self.up_revision_ids if self.is_upgrade else self.down_revision_ids 

859 ) 

860 

861 @property 

862 def up_revision(self): 

863 """Get :attr:`~.MigrationInfo.up_revision_id` as 

864 a :class:`.Revision`. 

865 

866 """ 

867 return self.revision_map.get_revision(self.up_revision_id) 

868 

869 @property 

870 def up_revisions(self): 

871 """Get :attr:`~.MigrationInfo.up_revision_ids` as a :class:`.Revision`. 

872 

873 .. versionadded:: 0.9.4 

874 

875 """ 

876 return self.revision_map.get_revisions(self.up_revision_ids) 

877 

878 @property 

879 def down_revisions(self): 

880 """Get :attr:`~.MigrationInfo.down_revision_ids` as a tuple of 

881 :class:`Revisions <.Revision>`.""" 

882 return self.revision_map.get_revisions(self.down_revision_ids) 

883 

884 @property 

885 def source_revisions(self): 

886 """Get :attr:`~MigrationInfo.source_revision_ids` as a tuple of 

887 :class:`Revisions <.Revision>`.""" 

888 return self.revision_map.get_revisions(self.source_revision_ids) 

889 

890 @property 

891 def destination_revisions(self): 

892 """Get :attr:`~MigrationInfo.destination_revision_ids` as a tuple of 

893 :class:`Revisions <.Revision>`.""" 

894 return self.revision_map.get_revisions(self.destination_revision_ids) 

895 

896 

897class MigrationStep(object): 

898 @property 

899 def name(self): 

900 return self.migration_fn.__name__ 

901 

902 @classmethod 

903 def upgrade_from_script(cls, revision_map, script): 

904 return RevisionStep(revision_map, script, True) 

905 

906 @classmethod 

907 def downgrade_from_script(cls, revision_map, script): 

908 return RevisionStep(revision_map, script, False) 

909 

910 @property 

911 def is_downgrade(self): 

912 return not self.is_upgrade 

913 

914 @property 

915 def short_log(self): 

916 return "%s %s -> %s" % ( 

917 self.name, 

918 util.format_as_comma(self.from_revisions_no_deps), 

919 util.format_as_comma(self.to_revisions_no_deps), 

920 ) 

921 

922 def __str__(self): 

923 if self.doc: 

924 return "%s %s -> %s, %s" % ( 

925 self.name, 

926 util.format_as_comma(self.from_revisions_no_deps), 

927 util.format_as_comma(self.to_revisions_no_deps), 

928 self.doc, 

929 ) 

930 else: 

931 return self.short_log 

932 

933 

934class RevisionStep(MigrationStep): 

935 def __init__(self, revision_map, revision, is_upgrade): 

936 self.revision_map = revision_map 

937 self.revision = revision 

938 self.is_upgrade = is_upgrade 

939 if is_upgrade: 

940 self.migration_fn = revision.module.upgrade 

941 else: 

942 self.migration_fn = revision.module.downgrade 

943 

944 def __repr__(self): 

945 return "RevisionStep(%r, is_upgrade=%r)" % ( 

946 self.revision.revision, 

947 self.is_upgrade, 

948 ) 

949 

950 def __eq__(self, other): 

951 return ( 

952 isinstance(other, RevisionStep) 

953 and other.revision == self.revision 

954 and self.is_upgrade == other.is_upgrade 

955 ) 

956 

957 @property 

958 def doc(self): 

959 return self.revision.doc 

960 

961 @property 

962 def from_revisions(self): 

963 if self.is_upgrade: 

964 return self.revision._all_down_revisions 

965 else: 

966 return (self.revision.revision,) 

967 

968 @property 

969 def from_revisions_no_deps(self): 

970 if self.is_upgrade: 

971 return self.revision._versioned_down_revisions 

972 else: 

973 return (self.revision.revision,) 

974 

975 @property 

976 def to_revisions(self): 

977 if self.is_upgrade: 

978 return (self.revision.revision,) 

979 else: 

980 return self.revision._all_down_revisions 

981 

982 @property 

983 def to_revisions_no_deps(self): 

984 if self.is_upgrade: 

985 return (self.revision.revision,) 

986 else: 

987 return self.revision._versioned_down_revisions 

988 

989 @property 

990 def _has_scalar_down_revision(self): 

991 return len(self.revision._all_down_revisions) == 1 

992 

993 def should_delete_branch(self, heads): 

994 """A delete is when we are a. in a downgrade and b. 

995 we are going to the "base" or we are going to a version that 

996 is implied as a dependency on another version that is remaining. 

997 

998 """ 

999 if not self.is_downgrade: 

1000 return False 

1001 

1002 if self.revision.revision not in heads: 

1003 return False 

1004 

1005 downrevs = self.revision._all_down_revisions 

1006 

1007 if not downrevs: 

1008 # is a base 

1009 return True 

1010 else: 

1011 # determine what the ultimate "to_revisions" for an 

1012 # unmerge would be. If there are none, then we're a delete. 

1013 to_revisions = self._unmerge_to_revisions(heads) 

1014 return not to_revisions 

1015 

1016 def merge_branch_idents(self, heads): 

1017 other_heads = set(heads).difference(self.from_revisions) 

1018 

1019 if other_heads: 

1020 ancestors = set( 

1021 r.revision 

1022 for r in self.revision_map._get_ancestor_nodes( 

1023 self.revision_map.get_revisions(other_heads), check=False 

1024 ) 

1025 ) 

1026 from_revisions = list( 

1027 set(self.from_revisions).difference(ancestors) 

1028 ) 

1029 else: 

1030 from_revisions = list(self.from_revisions) 

1031 

1032 return ( 

1033 # delete revs, update from rev, update to rev 

1034 list(from_revisions[0:-1]), 

1035 from_revisions[-1], 

1036 self.to_revisions[0], 

1037 ) 

1038 

1039 def _unmerge_to_revisions(self, heads): 

1040 other_heads = set(heads).difference([self.revision.revision]) 

1041 if other_heads: 

1042 ancestors = set( 

1043 r.revision 

1044 for r in self.revision_map._get_ancestor_nodes( 

1045 self.revision_map.get_revisions(other_heads), check=False 

1046 ) 

1047 ) 

1048 return list(set(self.to_revisions).difference(ancestors)) 

1049 else: 

1050 return self.to_revisions 

1051 

1052 def unmerge_branch_idents(self, heads): 

1053 to_revisions = self._unmerge_to_revisions(heads) 

1054 

1055 return ( 

1056 # update from rev, update to rev, insert revs 

1057 self.from_revisions[0], 

1058 to_revisions[-1], 

1059 to_revisions[0:-1], 

1060 ) 

1061 

1062 def should_create_branch(self, heads): 

1063 if not self.is_upgrade: 

1064 return False 

1065 

1066 downrevs = self.revision._all_down_revisions 

1067 

1068 if not downrevs: 

1069 # is a base 

1070 return True 

1071 else: 

1072 # none of our downrevs are present, so... 

1073 # we have to insert our version. This is true whether 

1074 # or not there is only one downrev, or multiple (in the latter 

1075 # case, we're a merge point.) 

1076 if not heads.intersection(downrevs): 

1077 return True 

1078 else: 

1079 return False 

1080 

1081 def should_merge_branches(self, heads): 

1082 if not self.is_upgrade: 

1083 return False 

1084 

1085 downrevs = self.revision._all_down_revisions 

1086 

1087 if len(downrevs) > 1 and len(heads.intersection(downrevs)) > 1: 

1088 return True 

1089 

1090 return False 

1091 

1092 def should_unmerge_branches(self, heads): 

1093 if not self.is_downgrade: 

1094 return False 

1095 

1096 downrevs = self.revision._all_down_revisions 

1097 

1098 if self.revision.revision in heads and len(downrevs) > 1: 

1099 return True 

1100 

1101 return False 

1102 

1103 def update_version_num(self, heads): 

1104 if not self._has_scalar_down_revision: 

1105 downrev = heads.intersection(self.revision._all_down_revisions) 

1106 assert ( 

1107 len(downrev) == 1 

1108 ), "Can't do an UPDATE because downrevision is ambiguous" 

1109 down_revision = list(downrev)[0] 

1110 else: 

1111 down_revision = self.revision._all_down_revisions[0] 

1112 

1113 if self.is_upgrade: 

1114 return down_revision, self.revision.revision 

1115 else: 

1116 return self.revision.revision, down_revision 

1117 

1118 @property 

1119 def delete_version_num(self): 

1120 return self.revision.revision 

1121 

1122 @property 

1123 def insert_version_num(self): 

1124 return self.revision.revision 

1125 

1126 @property 

1127 def info(self): 

1128 return MigrationInfo( 

1129 revision_map=self.revision_map, 

1130 up_revisions=self.revision.revision, 

1131 down_revisions=self.revision._all_down_revisions, 

1132 is_upgrade=self.is_upgrade, 

1133 is_stamp=False, 

1134 ) 

1135 

1136 

1137class StampStep(MigrationStep): 

1138 def __init__(self, from_, to_, is_upgrade, branch_move, revision_map=None): 

1139 self.from_ = util.to_tuple(from_, default=()) 

1140 self.to_ = util.to_tuple(to_, default=()) 

1141 self.is_upgrade = is_upgrade 

1142 self.branch_move = branch_move 

1143 self.migration_fn = self.stamp_revision 

1144 self.revision_map = revision_map 

1145 

1146 doc = None 

1147 

1148 def stamp_revision(self, **kw): 

1149 return None 

1150 

1151 def __eq__(self, other): 

1152 return ( 

1153 isinstance(other, StampStep) 

1154 and other.from_revisions == self.revisions 

1155 and other.to_revisions == self.to_revisions 

1156 and other.branch_move == self.branch_move 

1157 and self.is_upgrade == other.is_upgrade 

1158 ) 

1159 

1160 @property 

1161 def from_revisions(self): 

1162 return self.from_ 

1163 

1164 @property 

1165 def to_revisions(self): 

1166 return self.to_ 

1167 

1168 @property 

1169 def from_revisions_no_deps(self): 

1170 return self.from_ 

1171 

1172 @property 

1173 def to_revisions_no_deps(self): 

1174 return self.to_ 

1175 

1176 @property 

1177 def delete_version_num(self): 

1178 assert len(self.from_) == 1 

1179 return self.from_[0] 

1180 

1181 @property 

1182 def insert_version_num(self): 

1183 assert len(self.to_) == 1 

1184 return self.to_[0] 

1185 

1186 def update_version_num(self, heads): 

1187 assert len(self.from_) == 1 

1188 assert len(self.to_) == 1 

1189 return self.from_[0], self.to_[0] 

1190 

1191 def merge_branch_idents(self, heads): 

1192 return ( 

1193 # delete revs, update from rev, update to rev 

1194 list(self.from_[0:-1]), 

1195 self.from_[-1], 

1196 self.to_[0], 

1197 ) 

1198 

1199 def unmerge_branch_idents(self, heads): 

1200 return ( 

1201 # update from rev, update to rev, insert revs 

1202 self.from_[0], 

1203 self.to_[-1], 

1204 list(self.to_[0:-1]), 

1205 ) 

1206 

1207 def should_delete_branch(self, heads): 

1208 # TODO: we probably need to look for self.to_ inside of heads, 

1209 # in a similar manner as should_create_branch, however we have 

1210 # no tests for this yet (stamp downgrades w/ branches) 

1211 return self.is_downgrade and self.branch_move 

1212 

1213 def should_create_branch(self, heads): 

1214 return ( 

1215 self.is_upgrade 

1216 and (self.branch_move or set(self.from_).difference(heads)) 

1217 and set(self.to_).difference(heads) 

1218 ) 

1219 

1220 def should_merge_branches(self, heads): 

1221 return len(self.from_) > 1 

1222 

1223 def should_unmerge_branches(self, heads): 

1224 return len(self.to_) > 1 

1225 

1226 @property 

1227 def info(self): 

1228 up, down = ( 

1229 (self.to_, self.from_) 

1230 if self.is_upgrade 

1231 else (self.from_, self.to_) 

1232 ) 

1233 return MigrationInfo( 

1234 revision_map=self.revision_map, 

1235 up_revisions=up, 

1236 down_revisions=down, 

1237 is_upgrade=self.is_upgrade, 

1238 is_stamp=True, 

1239 )