74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
"""Registry for migration-specific prepare and verification hooks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Mapping
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING
|
|
|
|
if TYPE_CHECKING:
|
|
from tests.alembic.verifier import MigrationVerifier
|
|
|
|
|
|
BeforeUpgradeHook = Callable[["MigrationVerifier"], None]
|
|
AfterUpgradeHook = Callable[["MigrationVerifier"], None]
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class RevisionHooks:
|
|
"""Container for lifecycle hooks tied to a revision."""
|
|
|
|
before_upgrade: BeforeUpgradeHook | None = None
|
|
after_upgrade: AfterUpgradeHook | None = None
|
|
|
|
|
|
_REGISTRY: dict[str, RevisionHooks] = {}
|
|
|
|
|
|
def register_before_upgrade(
|
|
revision: str,
|
|
) -> Callable[[BeforeUpgradeHook], BeforeUpgradeHook]:
|
|
"""Register a callable to execute before upgrading to ``revision``."""
|
|
|
|
def decorator(func: BeforeUpgradeHook) -> BeforeUpgradeHook:
|
|
hooks = _REGISTRY.setdefault(revision, RevisionHooks())
|
|
if hooks.before_upgrade is not None:
|
|
raise ValueError(
|
|
f"before_upgrade hook already registered for revision {revision}"
|
|
)
|
|
hooks.before_upgrade = func
|
|
return func
|
|
|
|
return decorator
|
|
|
|
|
|
def register_after_upgrade(
|
|
revision: str,
|
|
) -> Callable[[AfterUpgradeHook], AfterUpgradeHook]:
|
|
"""Register a callable to execute after upgrading to ``revision``."""
|
|
|
|
def decorator(func: AfterUpgradeHook) -> AfterUpgradeHook:
|
|
hooks = _REGISTRY.setdefault(revision, RevisionHooks())
|
|
if hooks.after_upgrade is not None:
|
|
raise ValueError(
|
|
f"after_upgrade hook already registered for revision {revision}"
|
|
)
|
|
hooks.after_upgrade = func
|
|
return func
|
|
|
|
return decorator
|
|
|
|
|
|
def get_registered_hooks() -> Mapping[str, RevisionHooks]:
|
|
"""Return a snapshot of the registered revision hooks."""
|
|
|
|
return dict(_REGISTRY)
|
|
|
|
|
|
__all__ = [
|
|
"RevisionHooks",
|
|
"register_before_upgrade",
|
|
"register_after_upgrade",
|
|
"get_registered_hooks",
|
|
]
|