diff --git a/kanta/kantaimpl.py b/kanta/kantaimpl.py index 57c56f0..186f667 100644 --- a/kanta/kantaimpl.py +++ b/kanta/kantaimpl.py @@ -12,7 +12,7 @@ from typing import Any, Generic, TypeVar from kanta.callbacks import CallbackRegistry, InjectionContext from kanta.exceptions import DatabaseError, DataIntegrityError, ReplayError -from kanta.migrate import MigrationRegistry +from kanta.migrations import Migrations from kanta.persistence import PersistenceMixin from kanta.serialization import restore_data_in_place, struct_to_dict from kanta.serialization.base import replay @@ -28,18 +28,18 @@ class KantaImpl(PersistenceMixin, Generic[T]): def __init__(self, **kwargs: Any): self.data_type = kwargs.pop("type") self.data: T = kwargs.pop("data") - self.migrations = kwargs.pop("migrations", None) self._kanta = kwargs.pop("kanta", None) + migrations = kwargs.pop("migrations", None) self.ctx = SimpleNamespace() super().__init__(**kwargs) - self.migration_registry: MigrationRegistry | None = None - if self.migrations is not None: + self.migrations: Migrations | None = None + if migrations is not None: module = ( - importlib.import_module(self.migrations) - if isinstance(self.migrations, str) - else self.migrations + importlib.import_module(migrations) + if isinstance(migrations, str) + else migrations ) - self.migration_registry = MigrationRegistry.from_module(module) + self.migrations = Migrations.from_module(module) self.in_transaction = False self.transaction_snapshot: dict[str, Any] | None = None @@ -55,9 +55,7 @@ class KantaImpl(PersistenceMixin, Generic[T]): ) self.statedict = struct_to_dict(self.data, serializer=self.serializer) - self.version = ( - self.migration_registry.dbver if self.migration_registry is not None else 0 - ) + self.version = self.migrations.dbver if self.migrations is not None else 0 def add_bootstrap( self, @@ -141,10 +139,8 @@ class KantaImpl(PersistenceMixin, Generic[T]): cause_type=type(e).__name__, ) from e - if self.migration_registry is not None: - rr.version = self.migration_registry.apply( - rr.state, rr.version, self._kanta - ) + if self.migrations is not None: + rr.version = self.migrations.apply(rr.state, rr.version, self._kanta) self.statedict = copy.deepcopy(rr.state) self.data = restore_data_in_place( diff --git a/kanta/migrate.py b/kanta/migrations.py similarity index 80% rename from kanta/migrate.py rename to kanta/migrations.py index 33b820d..6e89882 100644 --- a/kanta/migrate.py +++ b/kanta/migrations.py @@ -14,29 +14,33 @@ from typing import Any _logger = logging.getLogger(__name__) +# Cache registries by imported module object so that many Kanta instances using +# the same migrations module do not re-scan it each time. +_module_registry_cache: dict[ModuleType, Migrations] = {} -class MigrationRegistry: + +class Migrations: """Registry of schema migration functions. Usage:: - registry = MigrationRegistry() + migrations = Migrations() - @registry.register + @migrations.register def migrate_v1(d: dict, kanta) -> None: d.setdefault("version", 1) kanta.ctx.note = "migrated" - @registry.register + @migrations.register def migrate_v2(d: dict) -> None: d.setdefault("version", 2) - new_version = registry.apply(state, current_version=0, kanta=kanta) + new_version = migrations.apply(state, current_version=0, kanta=kanta) Or load from a module:: - registry = MigrationRegistry.from_module("myapp.migrations") - new_version = registry.apply(state, current_version=0, kanta=kanta) + migrations = Migrations.from_module("myapp.migrations") + new_version = migrations.apply(state, current_version=0, kanta=kanta) """ def __init__(self) -> None: @@ -59,24 +63,30 @@ class MigrationRegistry: return fn @classmethod - def from_module(cls, module: str | ModuleType) -> MigrationRegistry: - """Create a registry by scanning a module for ``migrate_vN`` functions. + def from_module(cls, module: str | ModuleType) -> Migrations: + """Create or retrieve a cached registry by scanning a module. Args: module: A module name (string) or an imported module object. """ - reg = cls() if isinstance(module, str): mod = importlib.import_module(module) else: mod = module + try: + return _module_registry_cache[mod] + except KeyError: + pass + + reg = cls() for name in dir(mod): if name.startswith("migrate_v"): fn = getattr(mod, name) if callable(fn): version = reg._migration_version(fn) reg._migrations[version] = fn + _module_registry_cache[mod] = reg return reg @property diff --git a/tests/test_migrations.py b/tests/test_migrations.py index 1977497..a2b789d 100644 --- a/tests/test_migrations.py +++ b/tests/test_migrations.py @@ -1,6 +1,6 @@ from types import ModuleType, SimpleNamespace -from kanta.migrate import MigrationRegistry +from kanta.migrations import Migrations class _DummyKanta: @@ -9,7 +9,7 @@ class _DummyKanta: def test_register_and_apply(): - reg = MigrationRegistry() + reg = Migrations() kanta = _DummyKanta() @reg.register @@ -27,7 +27,7 @@ def test_register_and_apply(): def test_no_migrations_needed(): - reg = MigrationRegistry() + reg = Migrations() kanta = _DummyKanta() @reg.register @@ -52,7 +52,7 @@ def test_from_module(): mod.__dict__["migrate_v1"] = migrate_v1 mod.__dict__["migrate_v2"] = migrate_v2 - reg = MigrationRegistry.from_module(mod) + reg = Migrations.from_module(mod) assert reg.dbver == 2 state = {} @@ -62,7 +62,7 @@ def test_from_module(): def test_migrations_can_use_kanta_ctx(): - reg = MigrationRegistry() + reg = Migrations() kanta = _DummyKanta() @reg.register @@ -78,7 +78,7 @@ def test_migrations_can_use_kanta_ctx(): def test_migration_can_omit_kanta_argument(): - reg = MigrationRegistry() + reg = Migrations() kanta = _DummyKanta() @reg.register