ZpenrMessangerProject/venv/lib/python3.12/site-packages/alembic/testing/env.py

import importlib.machinery
import os
from pathlib import Path
import shutil
import textwrap

from sqlalchemy.testing import config
from sqlalchemy.testing import provision

from . import util as testing_util
from .. import command
from .. import script
from .. import util
from ..script import Script
from ..script import ScriptDirectory


def _get_staging_directory():
    if provision.FOLLOWER_IDENT:
        return f"scratch_{provision.FOLLOWER_IDENT}"
    else:
        return "scratch"


def staging_env(create=True, template="generic", sourceless=False):
    cfg = _testing_config()
    if create:
        path = _join_path(_get_staging_directory(), "scripts")
        assert not os.path.exists(path), (
            "staging directory %s already exists; poor cleanup?" % path
        )

        command.init(cfg, path, template=template)
        if sourceless:
            try:
                # do an import so that a .pyc/.pyo is generated.
                util.load_python_file(path, "env.py")
            except AttributeError:
                # we don't have the migration context set up yet
                # so running the .env py throws this exception.
                # theoretically we could be using py_compiler here to
                # generate .pyc/.pyo without importing but not really
                # worth it.
                pass
            assert sourceless in (
                "pep3147_envonly",
                "simple",
                "pep3147_everything",
            ), sourceless
            make_sourceless(
                _join_path(path, "env.py"),
                "pep3147" if "pep3147" in sourceless else "simple",
            )

    sc = script.ScriptDirectory.from_config(cfg)
    return sc


def clear_staging_env():
    from sqlalchemy.testing import engines

    engines.testing_reaper.close_all()
    shutil.rmtree(_get_staging_directory(), True)


def script_file_fixture(txt):
    dir_ = _join_path(_get_staging_directory(), "scripts")
    path = _join_path(dir_, "script.py.mako")
    with open(path, "w") as f:
        f.write(txt)


def env_file_fixture(txt):
    dir_ = _join_path(_get_staging_directory(), "scripts")
    txt = (
        """
from alembic import context

config = context.config
"""
        + txt
    )

    path = _join_path(dir_, "env.py")
    pyc_path = util.pyc_file_from_path(path)
    if pyc_path:
        os.unlink(pyc_path)

    with open(path, "w") as f:
        f.write(txt)


def _sqlite_file_db(tempname="foo.db", future=False, scope=None, **options):
    dir_ = _join_path(_get_staging_directory(), "scripts")
    url = "sqlite:///%s/%s" % (dir_, tempname)
    if scope:
        options["scope"] = scope
    return testing_util.testing_engine(url=url, future=future, options=options)


def _sqlite_testing_config(sourceless=False, future=False):
    dir_ = _join_path(_get_staging_directory(), "scripts")
    url = f"sqlite:///{dir_}/foo.db"

    sqlalchemy_future = future or ("future" in config.db.__class__.__module__)

    return _write_config_file(
        f"""
[alembic]
script_location = {dir_}
sqlalchemy.url = {url}
sourceless = {"true" if sourceless else "false"}
{"sqlalchemy.future = true" if sqlalchemy_future else ""}

[loggers]
keys = root,sqlalchemy

[handlers]
keys = console

[logger_root]
level = WARNING
handlers = console
qualname =

[logger_sqlalchemy]
level = DEBUG
handlers =
qualname = sqlalchemy.engine

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatters]
keys = generic

[formatter_generic]
format = %%(levelname)-5.5s [%%(name)s] %%(message)s
datefmt = %%H:%%M:%%S
    """
    )


def _multi_dir_testing_config(sourceless=False, extra_version_location=""):
    dir_ = _join_path(_get_staging_directory(), "scripts")
    sqlalchemy_future = "future" in config.db.__class__.__module__

    url = "sqlite:///%s/foo.db" % dir_

    return _write_config_file(
        f"""
[alembic]
script_location = {dir_}
sqlalchemy.url = {url}
sqlalchemy.future = {"true" if sqlalchemy_future else "false"}
sourceless = {"true" if sourceless else "false"}
path_separator = space
version_locations = %(here)s/model1/ %(here)s/model2/ %(here)s/model3/ \
{extra_version_location}

[loggers]
keys = root

[handlers]
keys = console

[logger_root]
level = WARNING
handlers = console
qualname =

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatters]
keys = generic

[formatter_generic]
format = %%(levelname)-5.5s [%%(name)s] %%(message)s
datefmt = %%H:%%M:%%S
    """
    )


def _no_sql_pyproject_config(dialect="postgresql", directives=""):
    """use a postgresql url with no host so that
    connections guaranteed to fail"""
    dir_ = _join_path(_get_staging_directory(), "scripts")

    return _write_toml_config(
        f"""
[tool.alembic]
script_location ="{dir_}"
{textwrap.dedent(directives)}

        """,
        f"""
[alembic]
sqlalchemy.url = {dialect}://

[loggers]
keys = root

[handlers]
keys = console

[logger_root]
level = WARNING
handlers = console
qualname =

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatters]
keys = generic

[formatter_generic]
format = %%(levelname)-5.5s [%%(name)s] %%(message)s
datefmt = %%H:%%M:%%S

""",
    )


def _no_sql_testing_config(dialect="postgresql", directives=""):
    """use a postgresql url with no host so that
    connections guaranteed to fail"""
    dir_ = _join_path(_get_staging_directory(), "scripts")
    return _write_config_file(
        f"""
[alembic]
script_location ={dir_}
sqlalchemy.url = {dialect}://
{directives}

[loggers]
keys = root

[handlers]
keys = console

[logger_root]
level = WARNING
handlers = console
qualname =

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatters]
keys = generic

[formatter_generic]
format = %%(levelname)-5.5s [%%(name)s] %%(message)s
datefmt = %%H:%%M:%%S

"""
    )


def _write_toml_config(tomltext, initext):
    cfg = _write_config_file(initext)
    with open(cfg.toml_file_name, "w") as f:
        f.write(tomltext)
    return cfg


def _write_config_file(text):
    cfg = _testing_config()
    with open(cfg.config_file_name, "w") as f:
        f.write(text)
    return cfg


def _testing_config():
    from alembic.config import Config

    if not os.access(_get_staging_directory(), os.F_OK):
        os.mkdir(_get_staging_directory())
    return Config(
        _join_path(_get_staging_directory(), "test_alembic.ini"),
        _join_path(_get_staging_directory(), "pyproject.toml"),
    )


def write_script(
    scriptdir, rev_id, content, encoding="ascii", sourceless=False
):
    old = scriptdir.revision_map.get_revision(rev_id)
    path = old.path

    content = textwrap.dedent(content)
    if encoding:
        content = content.encode(encoding)
    with open(path, "wb") as fp:
        fp.write(content)
    pyc_path = util.pyc_file_from_path(path)
    if pyc_path:
        os.unlink(pyc_path)
    script = Script._from_path(scriptdir, path)
    old = scriptdir.revision_map.get_revision(script.revision)
    if old.down_revision != script.down_revision:
        raise Exception("Can't change down_revision on a refresh operation.")
    scriptdir.revision_map.add_revision(script, _replace=True)

    if sourceless:
        make_sourceless(
            path, "pep3147" if sourceless == "pep3147_everything" else "simple"
        )


def make_sourceless(path, style):
    import py_compile

    py_compile.compile(path)

    if style == "simple":
        pyc_path = util.pyc_file_from_path(path)
        suffix = importlib.machinery.BYTECODE_SUFFIXES[0]
        filepath, ext = os.path.splitext(path)
        simple_pyc_path = filepath + suffix
        shutil.move(pyc_path, simple_pyc_path)
        pyc_path = simple_pyc_path
    else:
        assert style in ("pep3147", "simple")
        pyc_path = util.pyc_file_from_path(path)

    assert os.access(pyc_path, os.F_OK)

    os.unlink(path)


def three_rev_fixture(cfg):
    a = util.rev_id()
    b = util.rev_id()
    c = util.rev_id()

    script = ScriptDirectory.from_config(cfg)
    script.generate_revision(a, "revision a", refresh=True, head="base")
    write_script(
        script,
        a,
        f"""\
"Rev A"
revision = '{a}'
down_revision = None

from alembic import op


def upgrade():
    op.execute("CREATE STEP 1")


def downgrade():
    op.execute("DROP STEP 1")

""",
    )

    script.generate_revision(b, "revision b", refresh=True, head=a)
    write_script(
        script,
        b,
        f"""# coding: utf-8
"Rev B, méil, %3"
revision = '{b}'
down_revision = '{a}'

from alembic import op


def upgrade():
    op.execute("CREATE STEP 2")


def downgrade():
    op.execute("DROP STEP 2")

""",
        encoding="utf-8",
    )

    script.generate_revision(c, "revision c", refresh=True, head=b)
    write_script(
        script,
        c,
        f"""\
"Rev C"
revision = '{c}'
down_revision = '{b}'

from alembic import op


def upgrade():
    op.execute("CREATE STEP 3")


def downgrade():
    op.execute("DROP STEP 3")

""",
    )
    return a, b, c


def multi_heads_fixture(cfg, a, b, c):
    """Create a multiple head fixture from the three-revs fixture"""

    # a->b->c
    #     -> d -> e
    #     -> f
    d = util.rev_id()
    e = util.rev_id()
    f = util.rev_id()

    script = ScriptDirectory.from_config(cfg)
    script.generate_revision(
        d, "revision d from b", head=b, splice=True, refresh=True
    )
    write_script(
        script,
        d,
        f"""\
"Rev D"
revision = '{d}'
down_revision = '{b}'

from alembic import op


def upgrade():
    op.execute("CREATE STEP 4")


def downgrade():
    op.execute("DROP STEP 4")

""",
    )

    script.generate_revision(
        e, "revision e from d", head=d, splice=True, refresh=True
    )
    write_script(
        script,
        e,
        f"""\
"Rev E"
revision = '{e}'
down_revision = '{d}'

from alembic import op


def upgrade():
    op.execute("CREATE STEP 5")


def downgrade():
    op.execute("DROP STEP 5")

""",
    )

    script.generate_revision(
        f, "revision f from b", head=b, splice=True, refresh=True
    )
    write_script(
        script,
        f,
        f"""\
"Rev F"
revision = '{f}'
down_revision = '{b}'

from alembic import op


def upgrade():
    op.execute("CREATE STEP 6")


def downgrade():
    op.execute("DROP STEP 6")

""",
    )

    return d, e, f


def _multidb_testing_config(engines):
    """alembic.ini fixture to work exactly with the 'multidb' template"""

    dir_ = _join_path(_get_staging_directory(), "scripts")

    sqlalchemy_future = "future" in config.db.__class__.__module__

    databases = ", ".join(engines.keys())
    engines = "\n\n".join(
        f"[{key}]\nsqlalchemy.url = {value.url}"
        for key, value in engines.items()
    )

    return _write_config_file(
        f"""
[alembic]
script_location = {dir_}
sourceless = false
sqlalchemy.future = {"true" if sqlalchemy_future else "false"}
databases = {databases}

{engines}
[loggers]
keys = root

[handlers]
keys = console

[logger_root]
level = WARNING
handlers = console
qualname =

[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic

[formatters]
keys = generic

[formatter_generic]
format = %%(levelname)-5.5s [%%(name)s] %%(message)s
datefmt = %%H:%%M:%%S
    """
    )


def _join_path(base: str, *more: str):
    return str(Path(base).joinpath(*more).as_posix())