codekingpro/portable-devtools

1import argparse2from functools import wraps3import logging4import os5import sys6from flask import current_app, g7from alembic import __version__ as __alembic_version__8from alembic.config import Config as AlembicConfig9from alembic import command10from alembic.util import CommandError11 12alembic_version = tuple([int(v) for v in __alembic_version__.split('.')[0:3]])13log = logging.getLogger(__name__)14 15 16class _MigrateConfig(object):17    def __init__(self, migrate, db, **kwargs):18        self.migrate = migrate19        self.db = db20        self.directory = migrate.directory21        self.configure_args = kwargs22 23    @property24    def metadata(self):25        """26        Backwards compatibility, in old releases app.extensions['migrate']27        was set to db, and env.py accessed app.extensions['migrate'].metadata28        """29        return self.db.metadata30 31 32class Config(AlembicConfig):33    def __init__(self, *args, **kwargs):34        self.template_directory = kwargs.pop('template_directory', None)35        super().__init__(*args, **kwargs)36 37    def get_template_directory(self):38        if self.template_directory:39            return self.template_directory40        package_dir = os.path.abspath(os.path.dirname(__file__))41        return os.path.join(package_dir, 'templates')42 43 44class Migrate(object):45    def __init__(self, app=None, db=None, directory='migrations', command='db',46                 compare_type=True, render_as_batch=True, **kwargs):47        self.configure_callbacks = []48        self.db = db49        self.command = command50        self.directory = str(directory)51        self.alembic_ctx_kwargs = kwargs52        self.alembic_ctx_kwargs['compare_type'] = compare_type53        self.alembic_ctx_kwargs['render_as_batch'] = render_as_batch54        if app is not None and db is not None:55            self.init_app(app, db, directory)56 57    def init_app(self, app, db=None, directory=None, command=None,58                 compare_type=None, render_as_batch=None, **kwargs):59        self.db = db or self.db60        self.command = command or self.command61        self.directory = str(directory or self.directory)62        self.alembic_ctx_kwargs.update(kwargs)63        if compare_type is not None:64            self.alembic_ctx_kwargs['compare_type'] = compare_type65        if render_as_batch is not None:66            self.alembic_ctx_kwargs['render_as_batch'] = render_as_batch67        if not hasattr(app, 'extensions'):68            app.extensions = {}69        app.extensions['migrate'] = _MigrateConfig(70            self, self.db, **self.alembic_ctx_kwargs)71 72        from flask_migrate.cli import db as db_cli_group73        app.cli.add_command(db_cli_group, name=self.command)74 75    def configure(self, f):76        self.configure_callbacks.append(f)77        return f78 79    def call_configure_callbacks(self, config):80        for f in self.configure_callbacks:81            config = f(config)82        return config83 84    def get_config(self, directory=None, x_arg=None, opts=None):85        if directory is None:86            directory = self.directory87        directory = str(directory)88        config = Config(os.path.join(directory, 'alembic.ini'))89        config.set_main_option('script_location', directory)90        if config.cmd_opts is None:91            config.cmd_opts = argparse.Namespace()92        for opt in opts or []:93            setattr(config.cmd_opts, opt, True)94        if not hasattr(config.cmd_opts, 'x'):95            setattr(config.cmd_opts, 'x', [])96            for x in getattr(g, 'x_arg', []):97                config.cmd_opts.x.append(x)98            if x_arg is not None:99                if isinstance(x_arg, list) or isinstance(x_arg, tuple):100                    for x in x_arg:101                        config.cmd_opts.x.append(x)102                else:103                    config.cmd_opts.x.append(x_arg)104        return self.call_configure_callbacks(config)105 106 107def catch_errors(f):108    @wraps(f)109    def wrapped(*args, **kwargs):110        try:111            f(*args, **kwargs)112        except (CommandError, RuntimeError) as exc:113            log.error('Error: ' + str(exc))114            sys.exit(1)115    return wrapped116 117 118@catch_errors119def list_templates():120    """List available templates."""121    config = Config()122    config.print_stdout("Available templates:\n")123    for tempname in sorted(os.listdir(config.get_template_directory())):124        with open(125            os.path.join(config.get_template_directory(), tempname, "README")126        ) as readme:127            synopsis = next(readme).strip()128        config.print_stdout("%s - %s", tempname, synopsis)129 130 131@catch_errors132def init(directory=None, multidb=False, template=None, package=False):133    """Creates a new migration repository"""134    if directory is None:135        directory = current_app.extensions['migrate'].directory136    template_directory = None137    if template is not None and ('/' in template or '\\' in template):138        template_directory, template = os.path.split(template)139    config = Config(template_directory=template_directory)140    config.set_main_option('script_location', directory)141    config.config_file_name = os.path.join(directory, 'alembic.ini')142    config = current_app.extensions['migrate'].\143        migrate.call_configure_callbacks(config)144    if multidb and template is None:145        template = 'flask-multidb'146    elif template is None:147        template = 'flask'148    command.init(config, directory, template=template, package=package)149 150 151@catch_errors152def revision(directory=None, message=None, autogenerate=False, sql=False,153             head='head', splice=False, branch_label=None, version_path=None,154             rev_id=None):155    """Create a new revision file."""156    opts = ['autogenerate'] if autogenerate else None157    config = current_app.extensions['migrate'].migrate.get_config(158        directory, opts=opts)159    command.revision(config, message, autogenerate=autogenerate, sql=sql,160                     head=head, splice=splice, branch_label=branch_label,161                     version_path=version_path, rev_id=rev_id)162 163 164@catch_errors165def migrate(directory=None, message=None, sql=False, head='head', splice=False,166            branch_label=None, version_path=None, rev_id=None, x_arg=None):167    """Alias for 'revision --autogenerate'"""168    config = current_app.extensions['migrate'].migrate.get_config(169        directory, opts=['autogenerate'], x_arg=x_arg)170    command.revision(config, message, autogenerate=True, sql=sql,171                     head=head, splice=splice, branch_label=branch_label,172                     version_path=version_path, rev_id=rev_id)173 174 175@catch_errors176def edit(directory=None, revision='current'):177    """Edit current revision."""178    if alembic_version >= (0, 8, 0):179        config = current_app.extensions['migrate'].migrate.get_config(180            directory)181        command.edit(config, revision)182    else:183        raise RuntimeError('Alembic 0.8.0 or greater is required')184 185 186@catch_errors187def merge(directory=None, revisions='', message=None, branch_label=None,188          rev_id=None):189    """Merge two revisions together.  Creates a new migration file"""190    config = current_app.extensions['migrate'].migrate.get_config(directory)191    command.merge(config, revisions, message=message,192                  branch_label=branch_label, rev_id=rev_id)193 194 195@catch_errors196def upgrade(directory=None, revision='head', sql=False, tag=None, x_arg=None):197    """Upgrade to a later version"""198    config = current_app.extensions['migrate'].migrate.get_config(directory,199                                                                  x_arg=x_arg)200    command.upgrade(config, revision, sql=sql, tag=tag)201 202 203@catch_errors204def downgrade(directory=None, revision='-1', sql=False, tag=None, x_arg=None):205    """Revert to a previous version"""206    config = current_app.extensions['migrate'].migrate.get_config(directory,207                                                                  x_arg=x_arg)208    if sql and revision == '-1':209        revision = 'head:-1'210    command.downgrade(config, revision, sql=sql, tag=tag)211 212 213@catch_errors214def show(directory=None, revision='head'):215    """Show the revision denoted by the given symbol."""216    config = current_app.extensions['migrate'].migrate.get_config(directory)217    command.show(config, revision)218 219 220@catch_errors221def history(directory=None, rev_range=None, verbose=False,222            indicate_current=False):223    """List changeset scripts in chronological order."""224    config = current_app.extensions['migrate'].migrate.get_config(directory)225    if alembic_version >= (0, 9, 9):226        command.history(config, rev_range, verbose=verbose,227                        indicate_current=indicate_current)228    else:229        command.history(config, rev_range, verbose=verbose)230 231 232@catch_errors233def heads(directory=None, verbose=False, resolve_dependencies=False):234    """Show current available heads in the script directory"""235    config = current_app.extensions['migrate'].migrate.get_config(directory)236    command.heads(config, verbose=verbose,237                  resolve_dependencies=resolve_dependencies)238 239 240@catch_errors241def branches(directory=None, verbose=False):242    """Show current branch points"""243    config = current_app.extensions['migrate'].migrate.get_config(directory)244    command.branches(config, verbose=verbose)245 246 247@catch_errors248def current(directory=None, verbose=False):249    """Display the current revision for each database."""250    config = current_app.extensions['migrate'].migrate.get_config(directory)251    command.current(config, verbose=verbose)252 253 254@catch_errors255def stamp(directory=None, revision='head', sql=False, tag=None, purge=False):256    """'stamp' the revision table with the given revision; don't run any257    migrations"""258    config = current_app.extensions['migrate'].migrate.get_config(directory)259    command.stamp(config, revision, sql=sql, tag=tag, purge=purge)260 261 262@catch_errors263def check(directory=None):264    """Check if there are any new operations to migrate"""265    config = current_app.extensions['migrate'].migrate.get_config(directory)266    command.check(config)267