codekingpro/portable-devtools

1# mypy: allow-untyped-defs, allow-untyped-calls2 3from __future__ import annotations4 5import os6from typing import List7from typing import Optional8from typing import TYPE_CHECKING9from typing import Union10 11from . import autogenerate as autogen12from . import util13from .runtime.environment import EnvironmentContext14from .script import ScriptDirectory15 16if TYPE_CHECKING:17    from alembic.config import Config18    from alembic.script.base import Script19    from alembic.script.revision import _RevIdType20    from .runtime.environment import ProcessRevisionDirectiveFn21 22 23def list_templates(config: Config) -> None:24    """List available templates.25 26    :param config: a :class:`.Config` object.27 28    """29 30    config.print_stdout("Available templates:\n")31    for tempname in os.listdir(config.get_template_directory()):32        with open(33            os.path.join(config.get_template_directory(), tempname, "README")34        ) as readme:35            synopsis = next(readme).rstrip()36        config.print_stdout("%s - %s", tempname, synopsis)37 38    config.print_stdout("\nTemplates are used via the 'init' command, e.g.:")39    config.print_stdout("\n  alembic init --template generic ./scripts")40 41 42def init(43    config: Config,44    directory: str,45    template: str = "generic",46    package: bool = False,47) -> None:48    """Initialize a new scripts directory.49 50    :param config: a :class:`.Config` object.51 52    :param directory: string path of the target directory53 54    :param template: string name of the migration environment template to55     use.56 57    :param package: when True, write ``__init__.py`` files into the58     environment location as well as the versions/ location.59 60    """61 62    if os.access(directory, os.F_OK) and os.listdir(directory):63        raise util.CommandError(64            "Directory %s already exists and is not empty" % directory65        )66 67    template_dir = os.path.join(config.get_template_directory(), template)68    if not os.access(template_dir, os.F_OK):69        raise util.CommandError("No such template %r" % template)70 71    if not os.access(directory, os.F_OK):72        with util.status(73            f"Creating directory {os.path.abspath(directory)!r}",74            **config.messaging_opts,75        ):76            os.makedirs(directory)77 78    versions = os.path.join(directory, "versions")79    with util.status(80        f"Creating directory {os.path.abspath(versions)!r}",81        **config.messaging_opts,82    ):83        os.makedirs(versions)84 85    script = ScriptDirectory(directory)86 87    config_file: str | None = None88    for file_ in os.listdir(template_dir):89        file_path = os.path.join(template_dir, file_)90        if file_ == "alembic.ini.mako":91            assert config.config_file_name is not None92            config_file = os.path.abspath(config.config_file_name)93            if os.access(config_file, os.F_OK):94                util.msg(95                    f"File {config_file!r} already exists, skipping",96                    **config.messaging_opts,97                )98            else:99                script._generate_template(100                    file_path, config_file, script_location=directory101                )102        elif os.path.isfile(file_path):103            output_file = os.path.join(directory, file_)104            script._copy_file(file_path, output_file)105 106    if package:107        for path in [108            os.path.join(os.path.abspath(directory), "__init__.py"),109            os.path.join(os.path.abspath(versions), "__init__.py"),110        ]:111            with util.status(f"Adding {path!r}", **config.messaging_opts):112                with open(path, "w"):113                    pass114 115    assert config_file is not None116    util.msg(117        "Please edit configuration/connection/logging "118        f"settings in {config_file!r} before proceeding.",119        **config.messaging_opts,120    )121 122 123def revision(124    config: Config,125    message: Optional[str] = None,126    autogenerate: bool = False,127    sql: bool = False,128    head: str = "head",129    splice: bool = False,130    branch_label: Optional[_RevIdType] = None,131    version_path: Optional[str] = None,132    rev_id: Optional[str] = None,133    depends_on: Optional[str] = None,134    process_revision_directives: Optional[ProcessRevisionDirectiveFn] = None,135) -> Union[Optional[Script], List[Optional[Script]]]:136    """Create a new revision file.137 138    :param config: a :class:`.Config` object.139 140    :param message: string message to apply to the revision; this is the141     ``-m`` option to ``alembic revision``.142 143    :param autogenerate: whether or not to autogenerate the script from144     the database; this is the ``--autogenerate`` option to145     ``alembic revision``.146 147    :param sql: whether to dump the script out as a SQL string; when specified,148     the script is dumped to stdout.  This is the ``--sql`` option to149     ``alembic revision``.150 151    :param head: head revision to build the new revision upon as a parent;152     this is the ``--head`` option to ``alembic revision``.153 154    :param splice: whether or not the new revision should be made into a155     new head of its own; is required when the given ``head`` is not itself156     a head.  This is the ``--splice`` option to ``alembic revision``.157 158    :param branch_label: string label to apply to the branch; this is the159     ``--branch-label`` option to ``alembic revision``.160 161    :param version_path: string symbol identifying a specific version path162     from the configuration; this is the ``--version-path`` option to163     ``alembic revision``.164 165    :param rev_id: optional revision identifier to use instead of having166     one generated; this is the ``--rev-id`` option to ``alembic revision``.167 168    :param depends_on: optional list of "depends on" identifiers; this is the169     ``--depends-on`` option to ``alembic revision``.170 171    :param process_revision_directives: this is a callable that takes the172     same form as the callable described at173     :paramref:`.EnvironmentContext.configure.process_revision_directives`;174     will be applied to the structure generated by the revision process175     where it can be altered programmatically.   Note that unlike all176     the other parameters, this option is only available via programmatic177     use of :func:`.command.revision`178 179    """180 181    script_directory = ScriptDirectory.from_config(config)182 183    command_args = dict(184        message=message,185        autogenerate=autogenerate,186        sql=sql,187        head=head,188        splice=splice,189        branch_label=branch_label,190        version_path=version_path,191        rev_id=rev_id,192        depends_on=depends_on,193    )194    revision_context = autogen.RevisionContext(195        config,196        script_directory,197        command_args,198        process_revision_directives=process_revision_directives,199    )200 201    environment = util.asbool(config.get_main_option("revision_environment"))202 203    if autogenerate:204        environment = True205 206        if sql:207            raise util.CommandError(208                "Using --sql with --autogenerate does not make any sense"209            )210 211        def retrieve_migrations(rev, context):212            revision_context.run_autogenerate(rev, context)213            return []214 215    elif environment:216 217        def retrieve_migrations(rev, context):218            revision_context.run_no_autogenerate(rev, context)219            return []220 221    elif sql:222        raise util.CommandError(223            "Using --sql with the revision command when "224            "revision_environment is not configured does not make any sense"225        )226 227    if environment:228        with EnvironmentContext(229            config,230            script_directory,231            fn=retrieve_migrations,232            as_sql=sql,233            template_args=revision_context.template_args,234            revision_context=revision_context,235        ):236            script_directory.run_env()237 238        # the revision_context now has MigrationScript structure(s) present.239        # these could theoretically be further processed / rewritten *here*,240        # in addition to the hooks present within each run_migrations() call,241        # or at the end of env.py run_migrations_online().242 243    scripts = [script for script in revision_context.generate_scripts()]244    if len(scripts) == 1:245        return scripts[0]246    else:247        return scripts248 249 250def check(config: "Config") -> None:251    """Check if revision command with autogenerate has pending upgrade ops.252 253    :param config: a :class:`.Config` object.254 255    .. versionadded:: 1.9.0256 257    """258 259    script_directory = ScriptDirectory.from_config(config)260 261    command_args = dict(262        message=None,263        autogenerate=True,264        sql=False,265        head="head",266        splice=False,267        branch_label=None,268        version_path=None,269        rev_id=None,270        depends_on=None,271    )272    revision_context = autogen.RevisionContext(273        config,274        script_directory,275        command_args,276    )277 278    def retrieve_migrations(rev, context):279        revision_context.run_autogenerate(rev, context)280        return []281 282    with EnvironmentContext(283        config,284        script_directory,285        fn=retrieve_migrations,286        as_sql=False,287        template_args=revision_context.template_args,288        revision_context=revision_context,289    ):290        script_directory.run_env()291 292    # the revision_context now has MigrationScript structure(s) present.293 294    migration_script = revision_context.generated_revisions[-1]295    diffs = []296    for upgrade_ops in migration_script.upgrade_ops_list:297        diffs.extend(upgrade_ops.as_diffs())298 299    if diffs:300        raise util.AutogenerateDiffsDetected(301            f"New upgrade operations detected: {diffs}"302        )303    else:304        config.print_stdout("No new upgrade operations detected.")305 306 307def merge(308    config: Config,309    revisions: _RevIdType,310    message: Optional[str] = None,311    branch_label: Optional[_RevIdType] = None,312    rev_id: Optional[str] = None,313) -> Optional[Script]:314    """Merge two revisions together.  Creates a new migration file.315 316    :param config: a :class:`.Config` instance317 318    :param message: string message to apply to the revision319 320    :param branch_label: string label name to apply to the new revision321 322    :param rev_id: hardcoded revision identifier instead of generating a new323     one.324 325    .. seealso::326 327        :ref:`branches`328 329    """330 331    script = ScriptDirectory.from_config(config)332    template_args = {333        "config": config  # Let templates use config for334        # e.g. multiple databases335    }336 337    environment = util.asbool(config.get_main_option("revision_environment"))338 339    if environment:340 341        def nothing(rev, context):342            return []343 344        with EnvironmentContext(345            config,346            script,347            fn=nothing,348            as_sql=False,349            template_args=template_args,350        ):351            script.run_env()352 353    return script.generate_revision(354        rev_id or util.rev_id(),355        message,356        refresh=True,357        head=revisions,358        branch_labels=branch_label,359        **template_args,  # type:ignore[arg-type]360    )361 362 363def upgrade(364    config: Config,365    revision: str,366    sql: bool = False,367    tag: Optional[str] = None,368) -> None:369    """Upgrade to a later version.370 371    :param config: a :class:`.Config` instance.372 373    :param revision: string revision target or range for --sql mode374 375    :param sql: if True, use ``--sql`` mode376 377    :param tag: an arbitrary "tag" that can be intercepted by custom378     ``env.py`` scripts via the :meth:`.EnvironmentContext.get_tag_argument`379     method.380 381    """382 383    script = ScriptDirectory.from_config(config)384 385    starting_rev = None386    if ":" in revision:387        if not sql:388            raise util.CommandError("Range revision not allowed")389        starting_rev, revision = revision.split(":", 2)390 391    def upgrade(rev, context):392        return script._upgrade_revs(revision, rev)393 394    with EnvironmentContext(395        config,396        script,397        fn=upgrade,398        as_sql=sql,399        starting_rev=starting_rev,400        destination_rev=revision,401        tag=tag,402    ):403        script.run_env()404 405 406def downgrade(407    config: Config,408    revision: str,409    sql: bool = False,410    tag: Optional[str] = None,411) -> None:412    """Revert to a previous version.413 414    :param config: a :class:`.Config` instance.415 416    :param revision: string revision target or range for --sql mode417 418    :param sql: if True, use ``--sql`` mode419 420    :param tag: an arbitrary "tag" that can be intercepted by custom421     ``env.py`` scripts via the :meth:`.EnvironmentContext.get_tag_argument`422     method.423 424    """425 426    script = ScriptDirectory.from_config(config)427    starting_rev = None428    if ":" in revision:429        if not sql:430            raise util.CommandError("Range revision not allowed")431        starting_rev, revision = revision.split(":", 2)432    elif sql:433        raise util.CommandError(434            "downgrade with --sql requires <fromrev>:<torev>"435        )436 437    def downgrade(rev, context):438        return script._downgrade_revs(revision, rev)439 440    with EnvironmentContext(441        config,442        script,443        fn=downgrade,444        as_sql=sql,445        starting_rev=starting_rev,446        destination_rev=revision,447        tag=tag,448    ):449        script.run_env()450 451 452def show(config, rev):453    """Show the revision(s) denoted by the given symbol.454 455    :param config: a :class:`.Config` instance.456 457    :param revision: string revision target458 459    """460 461    script = ScriptDirectory.from_config(config)462 463    if rev == "current":464 465        def show_current(rev, context):466            for sc in script.get_revisions(rev):467                config.print_stdout(sc.log_entry)468            return []469 470        with EnvironmentContext(config, script, fn=show_current):471            script.run_env()472    else:473        for sc in script.get_revisions(rev):474            config.print_stdout(sc.log_entry)475 476 477def history(478    config: Config,479    rev_range: Optional[str] = None,480    verbose: bool = False,481    indicate_current: bool = False,482) -> None:483    """List changeset scripts in chronological order.484 485    :param config: a :class:`.Config` instance.486 487    :param rev_range: string revision range488 489    :param verbose: output in verbose mode.490 491    :param indicate_current: indicate current revision.492 493    """494    base: Optional[str]495    head: Optional[str]496    script = ScriptDirectory.from_config(config)497    if rev_range is not None:498        if ":" not in rev_range:499            raise util.CommandError(500                "History range requires [start]:[end], " "[start]:, or :[end]"501            )502        base, head = rev_range.strip().split(":")503    else:504        base = head = None505 506    environment = (507        util.asbool(config.get_main_option("revision_environment"))508        or indicate_current509    )510 511    def _display_history(config, script, base, head, currents=()):512        for sc in script.walk_revisions(513            base=base or "base", head=head or "heads"514        ):515            if indicate_current:516                sc._db_current_indicator = sc.revision in currents517 518            config.print_stdout(519                sc.cmd_format(520                    verbose=verbose,521                    include_branches=True,522                    include_doc=True,523                    include_parents=True,524                )525            )526 527    def _display_history_w_current(config, script, base, head):528        def _display_current_history(rev, context):529            if head == "current":530                _display_history(config, script, base, rev, rev)531            elif base == "current":532                _display_history(config, script, rev, head, rev)533            else:534                _display_history(config, script, base, head, rev)535            return []536 537        with EnvironmentContext(config, script, fn=_display_current_history):538            script.run_env()539 540    if base == "current" or head == "current" or environment:541        _display_history_w_current(config, script, base, head)542    else:543        _display_history(config, script, base, head)544 545 546def heads(config, verbose=False, resolve_dependencies=False):547    """Show current available heads in the script directory.548 549    :param config: a :class:`.Config` instance.550 551    :param verbose: output in verbose mode.552 553    :param resolve_dependencies: treat dependency version as down revisions.554 555    """556 557    script = ScriptDirectory.from_config(config)558    if resolve_dependencies:559        heads = script.get_revisions("heads")560    else:561        heads = script.get_revisions(script.get_heads())562 563    for rev in heads:564        config.print_stdout(565            rev.cmd_format(566                verbose, include_branches=True, tree_indicators=False567            )568        )569 570 571def branches(config, verbose=False):572    """Show current branch points.573 574    :param config: a :class:`.Config` instance.575 576    :param verbose: output in verbose mode.577 578    """579    script = ScriptDirectory.from_config(config)580    for sc in script.walk_revisions():581        if sc.is_branch_point:582            config.print_stdout(583                "%s\n%s\n",584                sc.cmd_format(verbose, include_branches=True),585                "\n".join(586                    "%s -> %s"587                    % (588                        " " * len(str(sc.revision)),589                        rev_obj.cmd_format(590                            False, include_branches=True, include_doc=verbose591                        ),592                    )593                    for rev_obj in (594                        script.get_revision(rev) for rev in sc.nextrev595                    )596                ),597            )598 599 600def current(config: Config, verbose: bool = False) -> None:601    """Display the current revision for a database.602 603    :param config: a :class:`.Config` instance.604 605    :param verbose: output in verbose mode.606 607    """608 609    script = ScriptDirectory.from_config(config)610 611    def display_version(rev, context):612        if verbose:613            config.print_stdout(614                "Current revision(s) for %s:",615                util.obfuscate_url_pw(context.connection.engine.url),616            )617        for rev in script.get_all_current(rev):618            config.print_stdout(rev.cmd_format(verbose))619 620        return []621 622    with EnvironmentContext(623        config, script, fn=display_version, dont_mutate=True624    ):625        script.run_env()626 627 628def stamp(629    config: Config,630    revision: _RevIdType,631    sql: bool = False,632    tag: Optional[str] = None,633    purge: bool = False,634) -> None:635    """'stamp' the revision table with the given revision; don't636    run any migrations.637 638    :param config: a :class:`.Config` instance.639 640    :param revision: target revision or list of revisions.   May be a list641     to indicate stamping of multiple branch heads.642 643     .. note:: this parameter is called "revisions" in the command line644        interface.645 646    :param sql: use ``--sql`` mode647 648    :param tag: an arbitrary "tag" that can be intercepted by custom649     ``env.py`` scripts via the :class:`.EnvironmentContext.get_tag_argument`650     method.651 652    :param purge: delete all entries in the version table before stamping.653 654    """655 656    script = ScriptDirectory.from_config(config)657 658    if sql:659        destination_revs = []660        starting_rev = None661        for _revision in util.to_list(revision):662            if ":" in _revision:663                srev, _revision = _revision.split(":", 2)664 665                if starting_rev != srev:666                    if starting_rev is None:667                        starting_rev = srev668                    else:669                        raise util.CommandError(670                            "Stamp operation with --sql only supports a "671                            "single starting revision at a time"672                        )673            destination_revs.append(_revision)674    else:675        destination_revs = util.to_list(revision)676 677    def do_stamp(rev, context):678        return script._stamp_revs(util.to_tuple(destination_revs), rev)679 680    with EnvironmentContext(681        config,682        script,683        fn=do_stamp,684        as_sql=sql,685        starting_rev=starting_rev if sql else None,686        destination_rev=util.to_tuple(destination_revs),687        tag=tag,688        purge=purge,689    ):690        script.run_env()691 692 693def edit(config: Config, rev: str) -> None:694    """Edit revision script(s) using $EDITOR.695 696    :param config: a :class:`.Config` instance.697 698    :param rev: target revision.699 700    """701 702    script = ScriptDirectory.from_config(config)703 704    if rev == "current":705 706        def edit_current(rev, context):707            if not rev:708                raise util.CommandError("No current revisions")709            for sc in script.get_revisions(rev):710                util.open_in_editor(sc.path)711            return []712 713        with EnvironmentContext(config, script, fn=edit_current):714            script.run_env()715    else:716        revs = script.get_revisions(rev)717        if not revs:718            raise util.CommandError(719                "No revision files indicated by symbol '%s'" % rev720            )721        for sc in revs:722            assert sc723            util.open_in_editor(sc.path)724 725 726def ensure_version(config: Config, sql: bool = False) -> None:727    """Create the alembic version table if it doesn't exist already .728 729    :param config: a :class:`.Config` instance.730 731    :param sql: use ``--sql`` mode732 733     .. versionadded:: 1.7.6734 735    """736 737    script = ScriptDirectory.from_config(config)738 739    def do_ensure_version(rev, context):740        context._ensure_version_table()741        return []742 743    with EnvironmentContext(744        config,745        script,746        fn=do_ensure_version,747        as_sql=sql,748    ):749        script.run_env()750