__init__.py 18 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360
  1. import os
  2. import argparse
  3. from flask import current_app
  4. from flask_script import Manager
  5. from alembic import __version__ as __alembic_version__
  6. from alembic.config import Config as AlembicConfig
  7. from alembic import command
  8. alembic_version = tuple([int(v) for v in __alembic_version__.split('.')[0:3]])
  9. class _MigrateConfig(object):
  10. def __init__(self, migrate, db, **kwargs):
  11. self.migrate = migrate
  12. self.db = db
  13. self.directory = migrate.directory
  14. self.configure_args = kwargs
  15. @property
  16. def metadata(self):
  17. """
  18. Backwards compatibility, in old releases app.extensions['migrate']
  19. was set to db, and env.py accessed app.extensions['migrate'].metadata
  20. """
  21. return self.db.metadata
  22. class Config(AlembicConfig):
  23. def get_template_directory(self):
  24. package_dir = os.path.abspath(os.path.dirname(__file__))
  25. return os.path.join(package_dir, 'templates')
  26. class Migrate(object):
  27. def __init__(self, app=None, db=None, directory='migrations', **kwargs):
  28. self.configure_callbacks = []
  29. self.db = db
  30. self.directory = directory
  31. if app is not None and db is not None:
  32. self.init_app(app, db, directory, **kwargs)
  33. def init_app(self, app, db=None, directory=None, **kwargs):
  34. self.db = db or self.db
  35. self.directory = directory or self.directory
  36. if not hasattr(app, 'extensions'):
  37. app.extensions = {}
  38. app.extensions['migrate'] = _MigrateConfig(self, self.db, **kwargs)
  39. def configure(self, f):
  40. self.configure_callbacks.append(f)
  41. return f
  42. def call_configure_callbacks(self, config):
  43. for f in self.configure_callbacks:
  44. config = f(config)
  45. return config
  46. def get_config(self, directory, x_arg=None, opts=None):
  47. if directory is None:
  48. directory = self.directory
  49. config = Config(os.path.join(directory, 'alembic.ini'))
  50. config.set_main_option('script_location', directory)
  51. if config.cmd_opts is None:
  52. config.cmd_opts = argparse.Namespace()
  53. for opt in opts or []:
  54. setattr(config.cmd_opts, opt, True)
  55. if x_arg is not None:
  56. if not getattr(config.cmd_opts, 'x', None):
  57. setattr(config.cmd_opts, 'x', [])
  58. if isinstance(x_arg, list) or isinstance(x_arg, tuple):
  59. for x in x_arg:
  60. config.cmd_opts.x.append(x)
  61. else:
  62. config.cmd_opts.x.append(x_arg)
  63. return self.call_configure_callbacks(config)
  64. MigrateCommand = Manager(usage='Perform database migrations')
  65. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  66. help=("migration script directory (default is "
  67. "'migrations')"))
  68. @MigrateCommand.option('--multidb', dest='multidb', action='store_true',
  69. default=False,
  70. help=("Multiple databases migraton (default is "
  71. "False)"))
  72. def init(directory=None, multidb=False):
  73. """Creates a new migration repository"""
  74. if directory is None:
  75. directory = current_app.extensions['migrate'].directory
  76. config = Config()
  77. config.set_main_option('script_location', directory)
  78. config.config_file_name = os.path.join(directory, 'alembic.ini')
  79. config = current_app.extensions['migrate'].\
  80. migrate.call_configure_callbacks(config)
  81. if multidb:
  82. command.init(config, directory, 'flask-multidb')
  83. else:
  84. command.init(config, directory, 'flask')
  85. @MigrateCommand.option('--rev-id', dest='rev_id', default=None,
  86. help=('Specify a hardcoded revision id instead of '
  87. 'generating one'))
  88. @MigrateCommand.option('--version-path', dest='version_path', default=None,
  89. help=('Specify specific path from config for version '
  90. 'file'))
  91. @MigrateCommand.option('--branch-label', dest='branch_label', default=None,
  92. help=('Specify a branch label to apply to the new '
  93. 'revision'))
  94. @MigrateCommand.option('--splice', dest='splice', action='store_true',
  95. default=False,
  96. help=('Allow a non-head revision as the "head" to '
  97. 'splice onto'))
  98. @MigrateCommand.option('--head', dest='head', default='head',
  99. help=('Specify head revision or <branchname>@head to '
  100. 'base new revision on'))
  101. @MigrateCommand.option('--sql', dest='sql', action='store_true', default=False,
  102. help=("Don't emit SQL to database - dump to standard "
  103. "output instead"))
  104. @MigrateCommand.option('--autogenerate', dest='autogenerate',
  105. action='store_true', default=False,
  106. help=('Populate revision script with andidate migration '
  107. 'operatons, based on comparison of database to '
  108. 'model'))
  109. @MigrateCommand.option('-m', '--message', dest='message', default=None,
  110. help='Revision message')
  111. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  112. help=("migration script directory (default is "
  113. "'migrations')"))
  114. def revision(directory=None, message=None, autogenerate=False, sql=False,
  115. head='head', splice=False, branch_label=None, version_path=None,
  116. rev_id=None):
  117. """Create a new revision file."""
  118. config = current_app.extensions['migrate'].migrate.get_config(directory)
  119. if alembic_version >= (0, 7, 0):
  120. command.revision(config, message, autogenerate=autogenerate, sql=sql,
  121. head=head, splice=splice, branch_label=branch_label,
  122. version_path=version_path, rev_id=rev_id)
  123. else:
  124. command.revision(config, message, autogenerate=autogenerate, sql=sql)
  125. @MigrateCommand.option('--rev-id', dest='rev_id', default=None,
  126. help=('Specify a hardcoded revision id instead of '
  127. 'generating one'))
  128. @MigrateCommand.option('--version-path', dest='version_path', default=None,
  129. help=('Specify specific path from config for version '
  130. 'file'))
  131. @MigrateCommand.option('--branch-label', dest='branch_label', default=None,
  132. help=('Specify a branch label to apply to the new '
  133. 'revision'))
  134. @MigrateCommand.option('--splice', dest='splice', action='store_true',
  135. default=False,
  136. help=('Allow a non-head revision as the "head" to '
  137. 'splice onto'))
  138. @MigrateCommand.option('--head', dest='head', default='head',
  139. help=('Specify head revision or <branchname>@head to '
  140. 'base new revision on'))
  141. @MigrateCommand.option('--sql', dest='sql', action='store_true', default=False,
  142. help=("Don't emit SQL to database - dump to standard "
  143. "output instead"))
  144. @MigrateCommand.option('-m', '--message', dest='message', default=None)
  145. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  146. help=("migration script directory (default is "
  147. "'migrations')"))
  148. def migrate(directory=None, message=None, sql=False, head='head', splice=False,
  149. branch_label=None, version_path=None, rev_id=None):
  150. """Alias for 'revision --autogenerate'"""
  151. config = current_app.extensions['migrate'].migrate.get_config(
  152. directory, opts=['autogenerate'])
  153. if alembic_version >= (0, 7, 0):
  154. command.revision(config, message, autogenerate=True, sql=sql,
  155. head=head, splice=splice, branch_label=branch_label,
  156. version_path=version_path, rev_id=rev_id)
  157. else:
  158. command.revision(config, message, autogenerate=True, sql=sql)
  159. @MigrateCommand.option('revision', nargs='?', default='head',
  160. help="revision identifier")
  161. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  162. help=("migration script directory (default is "
  163. "'migrations')"))
  164. def edit(directory=None, revision='current'):
  165. """Edit current revision."""
  166. if alembic_version >= (0, 8, 0):
  167. config = current_app.extensions['migrate'].migrate.get_config(
  168. directory)
  169. command.edit(config, revision)
  170. else:
  171. raise RuntimeError('Alembic 0.8.0 or greater is required')
  172. @MigrateCommand.option('--rev-id', dest='rev_id', default=None,
  173. help=('Specify a hardcoded revision id instead of '
  174. 'generating one'))
  175. @MigrateCommand.option('--branch-label', dest='branch_label', default=None,
  176. help=('Specify a branch label to apply to the new '
  177. 'revision'))
  178. @MigrateCommand.option('-m', '--message', dest='message', default=None)
  179. @MigrateCommand.option('revisions', nargs='+',
  180. help='one or more revisions, or "heads" for all heads')
  181. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  182. help=("migration script directory (default is "
  183. "'migrations')"))
  184. def merge(directory=None, revisions='', message=None, branch_label=None,
  185. rev_id=None):
  186. """Merge two revisions together. Creates a new migration file"""
  187. if alembic_version >= (0, 7, 0):
  188. config = current_app.extensions['migrate'].migrate.get_config(
  189. directory)
  190. command.merge(config, revisions, message=message,
  191. branch_label=branch_label, rev_id=rev_id)
  192. else:
  193. raise RuntimeError('Alembic 0.7.0 or greater is required')
  194. @MigrateCommand.option('--tag', dest='tag', default=None,
  195. help=("Arbitrary 'tag' name - can be used by custom "
  196. "env.py scripts"))
  197. @MigrateCommand.option('--sql', dest='sql', action='store_true', default=False,
  198. help=("Don't emit SQL to database - dump to standard "
  199. "output instead"))
  200. @MigrateCommand.option('revision', nargs='?', default='head',
  201. help="revision identifier")
  202. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  203. help=("migration script directory (default is "
  204. "'migrations')"))
  205. @MigrateCommand.option('-x', '--x-arg', dest='x_arg', default=None,
  206. help=("Additional arguments consumed by "
  207. "custom env.py scripts"))
  208. def upgrade(directory=None, revision='head', sql=False, tag=None, x_arg=None):
  209. """Upgrade to a later version"""
  210. config = current_app.extensions['migrate'].migrate.get_config(directory,
  211. x_arg=x_arg)
  212. command.upgrade(config, revision, sql=sql, tag=tag)
  213. @MigrateCommand.option('--tag', dest='tag', default=None,
  214. help=("Arbitrary 'tag' name - can be used by custom "
  215. "env.py scripts"))
  216. @MigrateCommand.option('--sql', dest='sql', action='store_true', default=False,
  217. help=("Don't emit SQL to database - dump to standard "
  218. "output instead"))
  219. @MigrateCommand.option('revision', nargs='?', default="-1",
  220. help="revision identifier")
  221. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  222. help=("migration script directory (default is "
  223. "'migrations')"))
  224. @MigrateCommand.option('-x', '--x-arg', dest='x_arg', default=None,
  225. help=("Additional arguments consumed by "
  226. "custom env.py scripts"))
  227. def downgrade(directory=None, revision='-1', sql=False, tag=None, x_arg=None):
  228. """Revert to a previous version"""
  229. config = current_app.extensions['migrate'].migrate.get_config(directory,
  230. x_arg=x_arg)
  231. if sql and revision == '-1':
  232. revision = 'head:-1'
  233. command.downgrade(config, revision, sql=sql, tag=tag)
  234. @MigrateCommand.option('revision', nargs='?', default="head",
  235. help="revision identifier")
  236. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  237. help=("migration script directory (default is "
  238. "'migrations')"))
  239. def show(directory=None, revision='head'):
  240. """Show the revision denoted by the given symbol."""
  241. if alembic_version >= (0, 7, 0):
  242. config = current_app.extensions['migrate'].migrate.get_config(
  243. directory)
  244. command.show(config, revision)
  245. else:
  246. raise RuntimeError('Alembic 0.7.0 or greater is required')
  247. @MigrateCommand.option('-v', '--verbose', dest='verbose', action='store_true',
  248. default=False, help='Use more verbose output')
  249. @MigrateCommand.option('-r', '--rev-range', dest='rev_range', default=None,
  250. help='Specify a revision range; format is [start]:[end]')
  251. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  252. help=("migration script directory (default is "
  253. "'migrations')"))
  254. def history(directory=None, rev_range=None, verbose=False):
  255. """List changeset scripts in chronological order."""
  256. config = current_app.extensions['migrate'].migrate.get_config(directory)
  257. if alembic_version >= (0, 7, 0):
  258. command.history(config, rev_range, verbose=verbose)
  259. else:
  260. command.history(config, rev_range)
  261. @MigrateCommand.option('--resolve-dependencies', dest='resolve_dependencies',
  262. action='store_true', default=False,
  263. help='Treat dependency versions as down revisions')
  264. @MigrateCommand.option('-v', '--verbose', dest='verbose', action='store_true',
  265. default=False, help='Use more verbose output')
  266. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  267. help=("migration script directory (default is "
  268. "'migrations')"))
  269. def heads(directory=None, verbose=False, resolve_dependencies=False):
  270. """Show current available heads in the script directory"""
  271. if alembic_version >= (0, 7, 0):
  272. config = current_app.extensions['migrate'].migrate.get_config(
  273. directory)
  274. command.heads(config, verbose=verbose,
  275. resolve_dependencies=resolve_dependencies)
  276. else:
  277. raise RuntimeError('Alembic 0.7.0 or greater is required')
  278. @MigrateCommand.option('-v', '--verbose', dest='verbose', action='store_true',
  279. default=False, help='Use more verbose output')
  280. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  281. help=("migration script directory (default is "
  282. "'migrations')"))
  283. def branches(directory=None, verbose=False):
  284. """Show current branch points"""
  285. config = current_app.extensions['migrate'].migrate.get_config(directory)
  286. if alembic_version >= (0, 7, 0):
  287. command.branches(config, verbose=verbose)
  288. else:
  289. command.branches(config)
  290. @MigrateCommand.option('--head-only', dest='head_only', action='store_true',
  291. default=False,
  292. help='Deprecated. Use --verbose for additional output')
  293. @MigrateCommand.option('-v', '--verbose', dest='verbose', action='store_true',
  294. default=False, help='Use more verbose output')
  295. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  296. help=("migration script directory (default is "
  297. "'migrations')"))
  298. def current(directory=None, verbose=False, head_only=False):
  299. """Display the current revision for each database."""
  300. config = current_app.extensions['migrate'].migrate.get_config(directory)
  301. if alembic_version >= (0, 7, 0):
  302. command.current(config, verbose=verbose, head_only=head_only)
  303. else:
  304. command.current(config)
  305. @MigrateCommand.option('--tag', dest='tag', default=None,
  306. help=("Arbitrary 'tag' name - can be used by custom "
  307. "env.py scripts"))
  308. @MigrateCommand.option('--sql', dest='sql', action='store_true', default=False,
  309. help=("Don't emit SQL to database - dump to standard "
  310. "output instead"))
  311. @MigrateCommand.option('revision', default=None, help="revision identifier")
  312. @MigrateCommand.option('-d', '--directory', dest='directory', default=None,
  313. help=("migration script directory (default is "
  314. "'migrations')"))
  315. def stamp(directory=None, revision='head', sql=False, tag=None):
  316. """'stamp' the revision table with the given revision; don't run any
  317. migrations"""
  318. config = current_app.extensions['migrate'].migrate.get_config(directory)
  319. command.stamp(config, revision, sql=sql, tag=tag)