diff --git a/alembic/script/__init__.py b/alembic/script/__init__.py index d78f3f1d..317db0e0 100644 --- a/alembic/script/__init__.py +++ b/alembic/script/__init__.py @@ -1,4 +1,5 @@ from .base import Script from .base import ScriptDirectory +from .revision import RevisionMap -__all__ = ["ScriptDirectory", "Script"] +__all__ = ["ScriptDirectory", "Script", "RevisionMap"] diff --git a/alembic/script/base.py b/alembic/script/base.py index f8417085..d0c13326 100644 --- a/alembic/script/base.py +++ b/alembic/script/base.py @@ -16,6 +16,7 @@ from typing import Sequence from typing import Set from typing import Tuple +from typing import Type from typing import TYPE_CHECKING from typing import Union @@ -86,6 +87,9 @@ def __init__( messaging_opts: MessagingOptions = cast( "MessagingOptions", util.EMPTY_DICT ), + revision_map_class: Optional[ + Type[revision.RevisionMap] + ] = None, ) -> None: self.dir = _preserving_path_as_str(dir) self.version_locations = [ @@ -95,7 +99,12 @@ def __init__( self.truncate_slug_length = truncate_slug_length or 40 self.sourceless = sourceless self.output_encoding = output_encoding - self.revision_map = revision.RevisionMap(self._load_revisions) + revision_map_cls = ( + revision_map_class + if revision_map_class is not None + else revision.RevisionMap + ) + self.revision_map = revision_map_cls(self._load_revisions) self.timezone = timezone self.hooks = hooks self.recursive_version_locations = recursive_version_locations @@ -183,6 +192,21 @@ def from_config(cls, config: Config) -> ScriptDirectory: sys.path[:0] = prepend_sys_path rvl = config.get_alembic_boolean_option("recursive_version_locations") + + revision_map_class: Optional[Type[revision.RevisionMap]] = None + rmc = config.get_alembic_option("revision_map_class") + if rmc is not None: + resolved = util.resolve_dotted_name(rmc) + if not ( + isinstance(resolved, type) + and issubclass(resolved, revision.RevisionMap) + ): + raise util.CommandError( + f"revision_map_class {rmc!r} must be a subclass of " + "alembic.script.revision.RevisionMap" + ) + revision_map_class = resolved + return ScriptDirectory( util.coerce_resource_to_filename(script_location), file_template=config.get_alembic_option( @@ -198,6 +222,7 @@ def from_config(cls, config: Config) -> ScriptDirectory: hooks=config.get_hooks_list(), recursive_version_locations=rvl, messaging_opts=config.messaging_opts, + revision_map_class=revision_map_class, ) @contextmanager diff --git a/alembic/templates/async/alembic.ini.mako b/alembic/templates/async/alembic.ini.mako index 02ccb0f6..be0f1851 100644 --- a/alembic/templates/async/alembic.ini.mako +++ b/alembic/templates/async/alembic.ini.mako @@ -79,6 +79,11 @@ path_separator = os # new in Alembic version 1.10 # recursive_version_locations = false +# specify a custom RevisionMap subclass for custom revision ordering logic. +# the value is a dotted Python path in "module:ClassName" format. +# the class must be a subclass of alembic.script.revision.RevisionMap. +# revision_map_class = mypackage.custom:CustomRevisionMap + # the output encoding used when revision files # are written from script.py.mako # output_encoding = utf-8 diff --git a/alembic/templates/generic/alembic.ini.mako b/alembic/templates/generic/alembic.ini.mako index 0127b2af..5484579a 100644 --- a/alembic/templates/generic/alembic.ini.mako +++ b/alembic/templates/generic/alembic.ini.mako @@ -79,6 +79,11 @@ path_separator = os # new in Alembic version 1.10 # recursive_version_locations = false +# specify a custom RevisionMap subclass for custom revision ordering logic. +# the value is a dotted Python path in "module:ClassName" format. +# the class must be a subclass of alembic.script.revision.RevisionMap. +# revision_map_class = mypackage.custom:CustomRevisionMap + # the output encoding used when revision files # are written from script.py.mako # output_encoding = utf-8 diff --git a/alembic/templates/multidb/alembic.ini.mako b/alembic/templates/multidb/alembic.ini.mako index 76846465..c50e066a 100644 --- a/alembic/templates/multidb/alembic.ini.mako +++ b/alembic/templates/multidb/alembic.ini.mako @@ -78,6 +78,11 @@ path_separator = os # new in Alembic version 1.10 # recursive_version_locations = false +# specify a custom RevisionMap subclass for custom revision ordering logic. +# the value is a dotted Python path in "module:ClassName" format. +# the class must be a subclass of alembic.script.revision.RevisionMap. +# revision_map_class = mypackage.custom:CustomRevisionMap + # the output encoding used when revision files # are written from script.py.mako # output_encoding = utf-8 diff --git a/alembic/templates/pyproject/pyproject.toml.mako b/alembic/templates/pyproject/pyproject.toml.mako index 7edd43b0..0802d27f 100644 --- a/alembic/templates/pyproject/pyproject.toml.mako +++ b/alembic/templates/pyproject/pyproject.toml.mako @@ -53,6 +53,11 @@ prepend_sys_path = [ # new in Alembic version 1.10 # recursive_version_locations = false +# specify a custom RevisionMap subclass for custom revision ordering logic. +# the value is a dotted Python path in "module:ClassName" format. +# the class must be a subclass of alembic.script.revision.RevisionMap. +# revision_map_class = "mypackage.custom:CustomRevisionMap" + # the output encoding used when revision files # are written from script.py.mako # output_encoding = "utf-8" diff --git a/alembic/templates/pyproject_async/pyproject.toml.mako b/alembic/templates/pyproject_async/pyproject.toml.mako index 7edd43b0..0802d27f 100644 --- a/alembic/templates/pyproject_async/pyproject.toml.mako +++ b/alembic/templates/pyproject_async/pyproject.toml.mako @@ -53,6 +53,11 @@ prepend_sys_path = [ # new in Alembic version 1.10 # recursive_version_locations = false +# specify a custom RevisionMap subclass for custom revision ordering logic. +# the value is a dotted Python path in "module:ClassName" format. +# the class must be a subclass of alembic.script.revision.RevisionMap. +# revision_map_class = "mypackage.custom:CustomRevisionMap" + # the output encoding used when revision files # are written from script.py.mako # output_encoding = "utf-8" diff --git a/alembic/util/__init__.py b/alembic/util/__init__.py index 8f3f685b..e6bfbdc0 100644 --- a/alembic/util/__init__.py +++ b/alembic/util/__init__.py @@ -29,5 +29,6 @@ from .pyfiles import coerce_resource_to_filename as coerce_resource_to_filename from .pyfiles import load_python_file as load_python_file from .pyfiles import pyc_file_from_path as pyc_file_from_path +from .pyfiles import resolve_dotted_name as resolve_dotted_name from .pyfiles import template_to_file as template_to_file from .sqla_compat import sqla_2 as sqla_2 diff --git a/alembic/util/pyfiles.py b/alembic/util/pyfiles.py index 135a42dc..eb0af83b 100644 --- a/alembic/util/pyfiles.py +++ b/alembic/util/pyfiles.py @@ -21,6 +21,41 @@ from .exc import CommandError +def resolve_dotted_name(dotted_name: str) -> Any: + """Resolve a dotted name string to a Python object. + + Accepts either ``"package.module:AttributeName"`` format (preferred, + following setuptools entry point conventions) or + ``"package.module.AttributeName"`` format. + + """ + if ":" in dotted_name: + module_part, _, attr_part = dotted_name.partition(":") + elif "." in dotted_name: + module_part, _, attr_part = dotted_name.rpartition(".") + else: + raise CommandError( + f"Could not resolve dotted name '{dotted_name}'; expected " + "format 'package.module:ClassName' or 'package.module.ClassName'" + ) + + try: + module = importlib.import_module(module_part) + except ImportError as ie: + raise CommandError( + f"Could not import module '{module_part}' " + f"from dotted name '{dotted_name}'" + ) from ie + + try: + return getattr(module, attr_part) + except AttributeError as ae: + raise CommandError( + f"Module '{module_part}' has no attribute '{attr_part}' " + f"(from dotted name '{dotted_name}')" + ) from ae + + def template_to_file( template_file: Union[str, os.PathLike[str]], dest: Union[str, os.PathLike[str]], diff --git a/docs/build/unreleased/revision_map_class.rst b/docs/build/unreleased/revision_map_class.rst new file mode 100644 index 00000000..21ba029d --- /dev/null +++ b/docs/build/unreleased/revision_map_class.rst @@ -0,0 +1,13 @@ +.. change:: + :tags: usecase, scripts + :tickets: + + Added a ``revision_map_class`` configuration option to + :class:`.ScriptDirectory`, allowing a custom subclass of + :class:`.RevisionMap` to be specified. This enables external packages to + implement custom revision ordering logic, such as ordering based on git + history, without changes to Alembic core. The option can be set in + ``alembic.ini`` or ``pyproject.toml`` using a dotted Python path in + ``"module:ClassName"`` format. :class:`.RevisionMap` is now also exported + from the ``alembic.script`` package for convenient subclassing. + diff --git a/tests/test_config.py b/tests/test_config.py index 0dc840f8..70d52e44 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -46,6 +46,41 @@ def tearDown(self): clear_staging_env() +class ResolveDottedNameTest(TestBase): + def test_colon_format(self): + result = util.resolve_dotted_name("collections:OrderedDict") + from collections import OrderedDict + + assert result is OrderedDict + + def test_dot_format(self): + result = util.resolve_dotted_name("collections.OrderedDict") + from collections import OrderedDict + + assert result is OrderedDict + + def test_invalid_no_separator(self): + with expect_raises_message( + util.CommandError, + r"Could not resolve dotted name", + ): + util.resolve_dotted_name("justaplainname") + + def test_bad_module(self): + with expect_raises_message( + util.CommandError, + r"Could not import module", + ): + util.resolve_dotted_name("nonexistent.module:Foo") + + def test_bad_attr(self): + with expect_raises_message( + util.CommandError, + r"has no attribute", + ): + util.resolve_dotted_name("collections:NoSuchThing") + + class ConfigTest(TestBase): def test_config_logging_with_file(self): buf = io.StringIO() @@ -752,6 +787,83 @@ def test_setting(self): eq_(script.output_encoding, "latin-1") +class RevisionMapClassTest(TestBase): + def setUp(self): + self.env = staging_env() + self.cfg = _no_sql_testing_config() + + def tearDown(self): + clear_staging_env() + + def test_default_revision_map_class(self): + from alembic.script.revision import RevisionMap + + script = ScriptDirectory.from_config(self.cfg) + assert type(script.revision_map) is RevisionMap + + def test_custom_class_via_constructor(self): + from alembic.script.revision import RevisionMap + + class CustomRevisionMap(RevisionMap): + pass + + script = ScriptDirectory( + self.cfg.get_main_option("script_location"), + revision_map_class=CustomRevisionMap, + ) + assert type(script.revision_map) is CustomRevisionMap + + def test_custom_class_from_config(self): + self.cfg.set_main_option( + "revision_map_class", + "alembic.script.revision:RevisionMap", + ) + script = ScriptDirectory.from_config(self.cfg) + from alembic.script.revision import RevisionMap + + assert type(script.revision_map) is RevisionMap + + def test_invalid_not_a_subclass(self): + self.cfg.set_main_option( + "revision_map_class", + "collections:OrderedDict", + ) + with expect_raises_message( + util.CommandError, + r"revision_map_class.*must be a subclass of " + r"alembic.script.revision.RevisionMap", + ): + ScriptDirectory.from_config(self.cfg) + + def test_invalid_unresolvable_module(self): + self.cfg.set_main_option( + "revision_map_class", + "nonexistent.module:FakeClass", + ) + with expect_raises_message( + util.CommandError, + r"Could not import module 'nonexistent.module'", + ): + ScriptDirectory.from_config(self.cfg) + + def test_invalid_unresolvable_attr(self): + self.cfg.set_main_option( + "revision_map_class", + "alembic.script.revision:NoSuchClass", + ) + with expect_raises_message( + util.CommandError, + r"has no attribute 'NoSuchClass'", + ): + ScriptDirectory.from_config(self.cfg) + + def test_revision_map_exported_from_script_package(self): + from alembic.script import RevisionMap + from alembic.script.revision import RevisionMap as DirectRevisionMap + + assert RevisionMap is DirectRevisionMap + + class CommandLineTest(TestBase): def test_register_command(self): cli = config.CommandLine()