alembic 1.15.2__py3-none-any.whl → 1.16.1__py3-none-any.whl

alembic/__init__.py

@@ -1,4 +1,4 @@
 from . import context
 from . import op
 
-__version__ = "1.15.2"
+__version__ = "1.16.1"

alembic/autogenerate/compare.py

@@ -24,6 +24,7 @@ from sqlalchemy import schema as sa_schema
 from sqlalchemy import text
 from sqlalchemy import types as sqltypes
 from sqlalchemy.sql import expression
+from sqlalchemy.sql.elements import conv
 from sqlalchemy.sql.schema import ForeignKeyConstraint
 from sqlalchemy.sql.schema import Index
 from sqlalchemy.sql.schema import UniqueConstraint
@@ -216,7 +217,7 @@ def _compare_tables(
                 (inspector),
                 # fmt: on
             )
-            inspector.reflect_table(t, include_columns=None)
+            _InspectorConv(inspector).reflect_table(t, include_columns=None)
         if autogen_context.run_object_filters(t, tname, "table", True, None):
             modify_table_ops = ops.ModifyTableOps(tname, [], schema=s)
 
@@ -246,7 +247,8 @@ def _compare_tables(
                 _compat_autogen_column_reflect(inspector),
                 # fmt: on
             )
-            inspector.reflect_table(t, include_columns=None)
+            _InspectorConv(inspector).reflect_table(t, include_columns=None)
+
         conn_column_info[(s, tname)] = t
 
     for s, tname in sorted(existing_tables, key=lambda x: (x[0] or "", x[1])):
@@ -438,6 +440,55 @@ def _compare_columns(
 _C = TypeVar("_C", bound=Union[UniqueConstraint, ForeignKeyConstraint, Index])
 
 
+class _InspectorConv:
+    __slots__ = ("inspector",)
+
+    def __init__(self, inspector):
+        self.inspector = inspector
+
+    def _apply_reflectinfo_conv(self, consts):
+        if not consts:
+            return consts
+        for const in consts:
+            if const["name"] is not None and not isinstance(
+                const["name"], conv
+            ):
+                const["name"] = conv(const["name"])
+        return consts
+
+    def _apply_constraint_conv(self, consts):
+        if not consts:
+            return consts
+        for const in consts:
+            if const.name is not None and not isinstance(const.name, conv):
+                const.name = conv(const.name)
+        return consts
+
+    def get_indexes(self, *args, **kw):
+        return self._apply_reflectinfo_conv(
+            self.inspector.get_indexes(*args, **kw)
+        )
+
+    def get_unique_constraints(self, *args, **kw):
+        return self._apply_reflectinfo_conv(
+            self.inspector.get_unique_constraints(*args, **kw)
+        )
+
+    def get_foreign_keys(self, *args, **kw):
+        return self._apply_reflectinfo_conv(
+            self.inspector.get_foreign_keys(*args, **kw)
+        )
+
+    def reflect_table(self, table, *, include_columns):
+        self.inspector.reflect_table(table, include_columns=include_columns)
+
+        # I had a cool version of this using _ReflectInfo, however that doesn't
+        # work in 1.4 and it's not public API in 2.x.  Then this is just a two
+        # liner.  So there's no competition...
+        self._apply_constraint_conv(table.constraints)
+        self._apply_constraint_conv(table.indexes)
+
+
 @comparators.dispatch_for("table")
 def _compare_indexes_and_uniques(
     autogen_context: AutogenContext,
@@ -473,9 +524,10 @@ def _compare_indexes_and_uniques(
     if conn_table is not None:
         # 1b. ... and from connection, if the table exists
         try:
-            conn_uniques = inspector.get_unique_constraints(  # type:ignore[assignment] # noqa
+            conn_uniques = _InspectorConv(inspector).get_unique_constraints(
                 tname, schema=schema
             )
+
             supports_unique_constraints = True
         except NotImplementedError:
             pass
@@ -498,7 +550,7 @@ def _compare_indexes_and_uniques(
                 if uq.get("duplicates_index"):
                     unique_constraints_duplicate_unique_indexes = True
         try:
-            conn_indexes = inspector.get_indexes(  # type:ignore[assignment]
+            conn_indexes = _InspectorConv(inspector).get_indexes(
                 tname, schema=schema
             )
         except NotImplementedError:
@@ -1178,7 +1230,9 @@ def _compare_foreign_keys(
 
     conn_fks_list = [
         fk
-        for fk in inspector.get_foreign_keys(tname, schema=schema)
+        for fk in _InspectorConv(inspector).get_foreign_keys(
+            tname, schema=schema
+        )
         if autogen_context.run_name_filters(
             fk["name"],
             "foreign_key_constraint",
@@ -1187,8 +1241,7 @@ def _compare_foreign_keys(
     ]
 
     conn_fks = {
-        _make_foreign_key(const, conn_table)  # type: ignore[arg-type]
-        for const in conn_fks_list
+        _make_foreign_key(const, conn_table) for const in conn_fks_list
     }
 
     impl = autogen_context.migration_context.impl

alembic/autogenerate/render.py

@@ -442,7 +442,7 @@ def _drop_constraint(
     name = _render_gen_name(autogen_context, op.constraint_name)
     schema = _ident(op.schema) if op.schema else None
     type_ = _ident(op.constraint_type) if op.constraint_type else None
-
+    if_exists = op.if_exists
     params_strs = []
     params_strs.append(repr(name))
     if not autogen_context._has_batch:
@@ -451,32 +451,47 @@ def _drop_constraint(
             params_strs.append(f"schema={schema!r}")
     if type_ is not None:
         params_strs.append(f"type_={type_!r}")
+    if if_exists is not None:
+        params_strs.append(f"if_exists={if_exists}")
 
     return f"{prefix}drop_constraint({', '.join(params_strs)})"
 
 
 @renderers.dispatch_for(ops.AddColumnOp)
 def _add_column(autogen_context: AutogenContext, op: ops.AddColumnOp) -> str:
-    schema, tname, column = op.schema, op.table_name, op.column
+    schema, tname, column, if_not_exists = (
+        op.schema,
+        op.table_name,
+        op.column,
+        op.if_not_exists,
+    )
     if autogen_context._has_batch:
         template = "%(prefix)sadd_column(%(column)s)"
     else:
         template = "%(prefix)sadd_column(%(tname)r, %(column)s"
         if schema:
             template += ", schema=%(schema)r"
+        if if_not_exists is not None:
+            template += ", if_not_exists=%(if_not_exists)r"
         template += ")"
     text = template % {
         "prefix": _alembic_autogenerate_prefix(autogen_context),
         "tname": tname,
         "column": _render_column(column, autogen_context),
         "schema": schema,
+        "if_not_exists": if_not_exists,
     }
     return text
 
 
 @renderers.dispatch_for(ops.DropColumnOp)
 def _drop_column(autogen_context: AutogenContext, op: ops.DropColumnOp) -> str:
-    schema, tname, column_name = op.schema, op.table_name, op.column_name
+    schema, tname, column_name, if_exists = (
+        op.schema,
+        op.table_name,
+        op.column_name,
+        op.if_exists,
+    )
 
     if autogen_context._has_batch:
         template = "%(prefix)sdrop_column(%(cname)r)"
@@ -484,6 +499,8 @@ def _drop_column(autogen_context: AutogenContext, op: ops.DropColumnOp) -> str:
         template = "%(prefix)sdrop_column(%(tname)r, %(cname)r"
         if schema:
             template += ", schema=%(schema)r"
+        if if_exists is not None:
+            template += ", if_exists=%(if_exists)r"
         template += ")"
 
     text = template % {
@@ -491,6 +508,7 @@ def _drop_column(autogen_context: AutogenContext, op: ops.DropColumnOp) -> str:
         "tname": _ident(tname),
         "cname": _ident(column_name),
         "schema": _ident(schema),
+        "if_exists": if_exists,
     }
     return text
 
@@ -1122,7 +1140,10 @@ def _execute_sql(autogen_context: AutogenContext, op: ops.ExecuteSQLOp) -> str:
             "Autogenerate rendering of SQL Expression language constructs "
             "not supported here; please use a plain SQL string"
         )
-    return "op.execute(%r)" % op.sqltext
+    return "{prefix}execute({sqltext!r})".format(
+        prefix=_alembic_autogenerate_prefix(autogen_context),
+        sqltext=op.sqltext,
+    )
 
 
 renderers = default_renderers.branch()

alembic/command.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import os
+import pathlib
 from typing import List
 from typing import Optional
 from typing import TYPE_CHECKING
@@ -12,6 +13,7 @@ from . import autogenerate as autogen
 from . import util
 from .runtime.environment import EnvironmentContext
 from .script import ScriptDirectory
+from .util import compat
 
 if TYPE_CHECKING:
     from alembic.config import Config
@@ -28,12 +30,10 @@ def list_templates(config: Config) -> None:
     """
 
     config.print_stdout("Available templates:\n")
-    for tempname in os.listdir(config.get_template_directory()):
-        with open(
-            os.path.join(config.get_template_directory(), tempname, "README")
-        ) as readme:
+    for tempname in config._get_template_path().iterdir():
+        with (tempname / "README").open() as readme:
             synopsis = next(readme).rstrip()
-        config.print_stdout("%s - %s", tempname, synopsis)
+        config.print_stdout("%s - %s", tempname.name, synopsis)
 
     config.print_stdout("\nTemplates are used via the 'init' command, e.g.:")
     config.print_stdout("\n  alembic init --template generic ./scripts")
@@ -59,65 +59,136 @@ def init(
 
     """
 
-    if os.access(directory, os.F_OK) and os.listdir(directory):
+    directory_path = pathlib.Path(directory)
+    if directory_path.exists() and list(directory_path.iterdir()):
         raise util.CommandError(
-            "Directory %s already exists and is not empty" % directory
+            "Directory %s already exists and is not empty" % directory_path
         )
 
-    template_dir = os.path.join(config.get_template_directory(), template)
-    if not os.access(template_dir, os.F_OK):
-        raise util.CommandError("No such template %r" % template)
+    template_path = config._get_template_path() / template
 
-    if not os.access(directory, os.F_OK):
+    if not template_path.exists():
+        raise util.CommandError(f"No such template {template_path}")
+
+    # left as os.access() to suit unit test mocking
+    if not os.access(directory_path, os.F_OK):
         with util.status(
-            f"Creating directory {os.path.abspath(directory)!r}",
+            f"Creating directory {directory_path.absolute()}",
             **config.messaging_opts,
         ):
-            os.makedirs(directory)
+            os.makedirs(directory_path)
 
-    versions = os.path.join(directory, "versions")
+    versions = directory_path / "versions"
     with util.status(
-        f"Creating directory {os.path.abspath(versions)!r}",
+        f"Creating directory {versions.absolute()}",
         **config.messaging_opts,
     ):
         os.makedirs(versions)
 
-    script = ScriptDirectory(directory)
+    if not directory_path.is_absolute():
+        # for non-absolute path, state config file in .ini / pyproject
+        # as relative to the %(here)s token, which is where the config
+        # file itself would be
+
+        if config._config_file_path is not None:
+            rel_dir = compat.path_relative_to(
+                directory_path.absolute(),
+                config._config_file_path.absolute().parent,
+                walk_up=True,
+            )
+            ini_script_location_directory = ("%(here)s" / rel_dir).as_posix()
+        if config._toml_file_path is not None:
+            rel_dir = compat.path_relative_to(
+                directory_path.absolute(),
+                config._toml_file_path.absolute().parent,
+                walk_up=True,
+            )
+            toml_script_location_directory = ("%(here)s" / rel_dir).as_posix()
+
+    else:
+        ini_script_location_directory = directory_path.as_posix()
+        toml_script_location_directory = directory_path.as_posix()
+
+    script = ScriptDirectory(directory_path)
 
-    config_file: str | None = None
-    for file_ in os.listdir(template_dir):
-        file_path = os.path.join(template_dir, file_)
+    has_toml = False
+
+    config_file: pathlib.Path | None = None
+
+    for file_path in template_path.iterdir():
+        file_ = file_path.name
         if file_ == "alembic.ini.mako":
             assert config.config_file_name is not None
-            config_file = os.path.abspath(config.config_file_name)
-            if os.access(config_file, os.F_OK):
+            config_file = pathlib.Path(config.config_file_name).absolute()
+            if config_file.exists():
                 util.msg(
-                    f"File {config_file!r} already exists, skipping",
+                    f"File {config_file} already exists, skipping",
                     **config.messaging_opts,
                 )
             else:
                 script._generate_template(
-                    file_path, config_file, script_location=directory
+                    file_path,
+                    config_file,
+                    script_location=ini_script_location_directory,
+                )
+        elif file_ == "pyproject.toml.mako":
+            has_toml = True
+            assert config._toml_file_path is not None
+            toml_path = config._toml_file_path.absolute()
+
+            if toml_path.exists():
+                # left as open() to suit unit test mocking
+                with open(toml_path, "rb") as f:
+                    toml_data = compat.tomllib.load(f)
+                    if "tool" in toml_data and "alembic" in toml_data["tool"]:
+
+                        util.msg(
+                            f"File {toml_path} already exists "
+                            "and already has a [tool.alembic] section, "
+                            "skipping",
+                        )
+                        continue
+                script._append_template(
+                    file_path,
+                    toml_path,
+                    script_location=toml_script_location_directory,
                 )
-        elif os.path.isfile(file_path):
-            output_file = os.path.join(directory, file_)
+            else:
+                script._generate_template(
+                    file_path,
+                    toml_path,
+                    script_location=toml_script_location_directory,
+                )
+
+        elif file_path.is_file():
+            output_file = directory_path / file_
             script._copy_file(file_path, output_file)
 
     if package:
         for path in [
-            os.path.join(os.path.abspath(directory), "__init__.py"),
-            os.path.join(os.path.abspath(versions), "__init__.py"),
+            directory_path.absolute() / "__init__.py",
+            versions.absolute() / "__init__.py",
         ]:
-            with util.status(f"Adding {path!r}", **config.messaging_opts):
+            with util.status(f"Adding {path!s}", **config.messaging_opts):
+                # left as open() to suit unit test mocking
                 with open(path, "w"):
                     pass
 
     assert config_file is not None
-    util.msg(
-        "Please edit configuration/connection/logging "
-        f"settings in {config_file!r} before proceeding.",
-        **config.messaging_opts,
-    )
+
+    if has_toml:
+        util.msg(
+            f"Please edit configuration settings in {toml_path} and "
+            "configuration/connection/logging "
+            f"settings in {config_file} before proceeding.",
+            **config.messaging_opts,
+        )
+    else:
+        util.msg(
+            "Please edit configuration/connection/logging "
+            f"settings in {config_file} before proceeding.",
+            **config.messaging_opts,
+        )
 
 
 def revision(
@@ -128,7 +199,7 @@ def revision(
     head: str = "head",
     splice: bool = False,
     branch_label: Optional[_RevIdType] = None,
-    version_path: Optional[str] = None,
+    version_path: Union[str, os.PathLike[str], None] = None,
     rev_id: Optional[str] = None,
     depends_on: Optional[str] = None,
     process_revision_directives: Optional[ProcessRevisionDirectiveFn] = None,
@@ -198,7 +269,9 @@ def revision(
         process_revision_directives=process_revision_directives,
     )
 
-    environment = util.asbool(config.get_main_option("revision_environment"))
+    environment = util.asbool(
+        config.get_alembic_option("revision_environment")
+    )
 
     if autogenerate:
         environment = True
@@ -338,7 +411,9 @@ def merge(
         # e.g. multiple databases
     }
 
-    environment = util.asbool(config.get_main_option("revision_environment"))
+    environment = util.asbool(
+        config.get_alembic_option("revision_environment")
+    )
 
     if environment:
 
@@ -511,7 +586,7 @@ def history(
         base = head = None
 
     environment = (
-        util.asbool(config.get_main_option("revision_environment"))
+        util.asbool(config.get_alembic_option("revision_environment"))
         or indicate_current
     )