diff --git a/kanta/kanta.py b/kanta/kanta.py index 9f1e80f..64c1b76 100644 --- a/kanta/kanta.py +++ b/kanta/kanta.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime from pathlib import Path -from types import ModuleType +from types import ModuleType, SimpleNamespace from typing import Any, Generic, TypeVar from kanta.kantaimpl import KantaImpl @@ -50,7 +50,6 @@ class Kanta(Generic[T]): *, type: type[T] | None = None, migrations: ModuleType | str | None = None, - migration_ctx: Any | None = None, serializer: Serializer | None = None, flush_interval: float = 0.1, ): @@ -61,7 +60,6 @@ class Kanta(Generic[T]): data: Caller-owned root msgspec.Struct state instance. type: Optional explicit root type. Defaults to ``type(data)``. migrations: Optional migrations module object or import path. - migration_ctx: Optional context object passed to migration functions. flush_interval: Background flush interval in seconds. serializer: Optional serializer implementation. @@ -78,7 +76,6 @@ class Kanta(Generic[T]): data=data, type=data_type, migrations=migrations, - migration_ctx=migration_ctx, flush_interval=flush_interval, kanta=self, ) @@ -127,6 +124,15 @@ class Kanta(Generic[T]): """ return self._impl.filename + @property + def ctx(self) -> SimpleNamespace: + """User-writable context namespace. + + Migration functions receive the ``Kanta`` instance and can read or + mutate ``kanta.ctx`` during migrations. + """ + return self._impl.ctx + @property def mtime(self) -> datetime | None: """Last modification time carried forward from change records. diff --git a/kanta/kantaimpl.py b/kanta/kantaimpl.py index eed0343..1471238 100644 --- a/kanta/kantaimpl.py +++ b/kanta/kantaimpl.py @@ -7,6 +7,7 @@ import copy import importlib import logging from datetime import UTC, datetime +from types import SimpleNamespace from typing import Any, Generic, TypeVar from kanta.callbacks import CallbackRegistry, InjectionContext @@ -28,8 +29,8 @@ class KantaImpl(PersistenceMixin, Generic[T]): self.data_type = kwargs.pop("type") self.data: T = kwargs.pop("data") self.migrations = kwargs.pop("migrations", None) - self.migration_ctx = kwargs.pop("migration_ctx", None) self._kanta = kwargs.pop("kanta", None) + self.ctx = SimpleNamespace() super().__init__(**kwargs) self.migration_registry: MigrationRegistry | None = None if self.migrations is not None: @@ -142,7 +143,7 @@ class KantaImpl(PersistenceMixin, Generic[T]): if self.migration_registry is not None: rr.version = self.migration_registry.apply( - rr.state, rr.version, self.migration_ctx + rr.state, rr.version, self._kanta ) self.statedict = copy.deepcopy(rr.state) diff --git a/kanta/migrate.py b/kanta/migrate.py index dccfc04..33b820d 100644 --- a/kanta/migrate.py +++ b/kanta/migrate.py @@ -7,24 +7,14 @@ or by prefix. Each runs exactly once based on the current version. from __future__ import annotations import importlib +import inspect import logging from types import ModuleType from typing import Any -import msgspec - _logger = logging.getLogger(__name__) -class MigrationCtx(msgspec.Struct, omit_defaults=True): - """Context passed to each migration function. - - Subclass or replace this with your own context type. - """ - - pass - - class MigrationRegistry: """Registry of schema migration functions. @@ -33,15 +23,20 @@ class MigrationRegistry: registry = MigrationRegistry() @registry.register - def migrate_v1(d: dict, ctx: MigrationCtx) -> None: + def migrate_v1(d: dict, kanta) -> None: d.setdefault("version", 1) + kanta.ctx.note = "migrated" - new_version = registry.apply(state, current_version=0) + @registry.register + def migrate_v2(d: dict) -> None: + d.setdefault("version", 2) + + new_version = registry.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) + new_version = registry.apply(state, current_version=0, kanta=kanta) """ def __init__(self) -> None: @@ -89,11 +84,21 @@ class MigrationRegistry: """Current schema version (= highest discovered migration, or 0).""" return max(self._migrations.keys(), default=0) + @staticmethod + def _call_migration(fn: Any, data_dict: dict[str, Any], kanta: Any) -> None: + """Call *fn* with the data dict and, if accepted, the Kanta instance.""" + try: + inspect.signature(fn).bind(data_dict, kanta) + except TypeError: + fn(data_dict) + else: + fn(data_dict, kanta) + def apply( self, data_dict: dict[str, Any], current_version: int, - ctx: MigrationCtx | None = None, + kanta: Any, *, silent: bool = False, ) -> int: @@ -109,7 +114,7 @@ class MigrationRegistry: f"Missing migration step migrate_v{next_version} " f"(highest discovered is v{self.dbver})" ) - fn(data_dict, ctx or MigrationCtx()) + self._call_migration(fn, data_dict, kanta) current_version = next_version if not silent: desc = (fn.__doc__ or fn.__name__).split("\n")[0].rstrip(".") diff --git a/tests/test_kanta_integration.py b/tests/test_kanta_integration.py index eff6179..613eb7b 100644 --- a/tests/test_kanta_integration.py +++ b/tests/test_kanta_integration.py @@ -404,7 +404,7 @@ async def test_migrations_from_module(tmp_path, format_config): mod = type(sys)("test_migrations") - def migrate_v1(d, ctx): + def migrate_v1(d, kanta): d["version"] = 1 mod.__dict__["migrate_v1"] = migrate_v1 @@ -514,7 +514,7 @@ async def test_migrations_from_module_path(tmp_path, format_config): module_name = "test_migrations_path" mod = type(sys)(module_name) - def migrate_v1(d, ctx): + def migrate_v1(d, kanta): d["counter"] = 2 mod.__dict__["migrate_v1"] = migrate_v1 diff --git a/tests/test_migrations.py b/tests/test_migrations.py index a944e7a..1977497 100644 --- a/tests/test_migrations.py +++ b/tests/test_migrations.py @@ -1,44 +1,52 @@ -from types import ModuleType +from types import ModuleType, SimpleNamespace from kanta.migrate import MigrationRegistry +class _DummyKanta: + def __init__(self): + self.ctx = SimpleNamespace() + + def test_register_and_apply(): reg = MigrationRegistry() + kanta = _DummyKanta() @reg.register - def migrate_v1(d, ctx): + def migrate_v1(d, kanta): d["version"] = 1 @reg.register - def migrate_v2(d, ctx): + def migrate_v2(d, kanta): d["version"] = 2 state = {} - new_ver = reg.apply(state, current_version=0, silent=True) + new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True) assert new_ver == 2 assert state["version"] == 2 def test_no_migrations_needed(): reg = MigrationRegistry() + kanta = _DummyKanta() @reg.register - def migrate_v1(d, ctx): + def migrate_v1(d, kanta): d["x"] = 1 state = {"x": 1} - new_ver = reg.apply(state, current_version=1, silent=True) + new_ver = reg.apply(state, current_version=1, kanta=kanta, silent=True) assert new_ver == 1 def test_from_module(): mod = ModuleType("fake_migrations") + kanta = _DummyKanta() - def migrate_v1(d, ctx): + def migrate_v1(d, kanta): d["v"] = 1 - def migrate_v2(d, ctx): + def migrate_v2(d, kanta): d["v"] = 2 mod.__dict__["migrate_v1"] = migrate_v1 @@ -48,6 +56,36 @@ def test_from_module(): assert reg.dbver == 2 state = {} - new_ver = reg.apply(state, current_version=0, silent=True) + new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True) assert new_ver == 2 assert state["v"] == 2 + + +def test_migrations_can_use_kanta_ctx(): + reg = MigrationRegistry() + kanta = _DummyKanta() + + @reg.register + def migrate_v1(d, kanta): + kanta.ctx.source = "migration" + d["source"] = kanta.ctx.source + + state = {} + new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True) + assert new_ver == 1 + assert state["source"] == "migration" + assert kanta.ctx.source == "migration" + + +def test_migration_can_omit_kanta_argument(): + reg = MigrationRegistry() + kanta = _DummyKanta() + + @reg.register + def migrate_v1(d): + d["x"] = 1 + + state = {} + new_ver = reg.apply(state, current_version=0, kanta=kanta, silent=True) + assert new_ver == 1 + assert state["x"] == 1 diff --git a/tests/test_readonly.py b/tests/test_readonly.py index b84ed8e..878b6e5 100644 --- a/tests/test_readonly.py +++ b/tests/test_readonly.py @@ -116,7 +116,7 @@ async def test_readonly_runs_migrations(tmp_path, format_config): format_config, ) - def migrate_v1(data, ctx): + def migrate_v1(data, kanta): data.setdefault("enabled", True) migrations = make_migrations_module("readonly_migrations", "migrate_v1", migrate_v1)