Add flag to migrate metadata on load (#605)

Author: mike0sv

mlem/api/migrations.py

@@ -30,16 +30,21 @@ def migrate(path: str, project: Optional[str] = None, recursive: bool = False):
         _migrate_one(loc)
 
 
-def _migrate_one(location: Location):
-    with location.open("r") as f:
-        payload = safe_load(f)
-
+def apply_migrations(payload: dict):
     changed = False
     for migration in _migrations:
         migrated = migration(payload)
         if migrated is not None:
             payload = migrated
             changed = True
+    return payload, changed
+
+
+def _migrate_one(location: Location):
+    with location.open("r") as f:
+        payload = safe_load(f)
+
+    payload, changed = apply_migrations(payload)
 
     if changed:
         echo(f"Migrated MLEM Object at {location}")

mlem/core/metadata.py

@@ -138,6 +138,7 @@ def load(
     rev: Optional[str] = None,
     batch_size: Optional[int] = None,
     follow_links: bool = True,
+    try_migrations: bool = False,
 ) -> Any:
     """Load python object saved by MLEM
 
@@ -158,6 +159,7 @@ def load(
         rev=rev,
         follow_links=follow_links,
         load_value=batch_size is None,
+        try_migrations=try_migrations,
     )
     if isinstance(meta, MlemData) and batch_size:
         return meta.read_batch(batch_size)
@@ -175,6 +177,7 @@ def load_meta(
     follow_links: bool = True,
     load_value: bool = False,
     fs: Optional[AbstractFileSystem] = None,
+    try_migrations: bool = False,
     *,
     force_type: Literal[None] = None,
 ) -> MlemObject:
@@ -189,6 +192,7 @@ def load_meta(
     follow_links: bool = True,
     load_value: bool = False,
     fs: Optional[AbstractFileSystem] = None,
+    try_migrations: bool = False,
     *,
     force_type: Optional[Type[T]] = None,
 ) -> T:
@@ -203,6 +207,7 @@ def load_meta(
     follow_links: bool = True,
     load_value: bool = False,
     fs: Optional[AbstractFileSystem] = None,
+    try_migrations: bool = False,
     *,
     force_type: Optional[Type[T]] = None,
 ) -> T:
@@ -216,6 +221,7 @@ def load_meta(
             actual object link points to. Defaults to True.
         load_value: Load actual python object incorporated in MlemObject. Defaults to False.
         fs: filesystem to load from. If not provided, will be inferred from path
+        try_migrations: If loading older versions of metadata, try to apply migrations
         force_type: type of meta to be loaded. Defaults to MlemObject (any mlem meta)
     Returns:
         MlemObject: Saved MlemObject
@@ -232,6 +238,7 @@ def load_meta(
     meta = cls.read(
         location=find_meta_location(location),
         follow_links=follow_links,
+        try_migrations=try_migrations,
     )
     log_meta_params(meta, add_object_type=True)
     if load_value:
@@ -274,14 +281,21 @@ def find_meta_location(location: Location) -> Location:
 
 
 def list_objects(
-    path: str = ".", fs: Optional[AbstractFileSystem] = None, recursive=True
+    path: str = ".",
+    fs: Optional[AbstractFileSystem] = None,
+    recursive=True,
+    try_migrations=False,
 ) -> Dict[Type[MlemObject], List[MlemObject]]:
     loc = Location.resolve(path, fs=fs)
     result = defaultdict(list)
     postfix = f"/**{MLEM_EXT}" if recursive else f"/*{MLEM_EXT}"
     for filepath in loc.fs.glob(loc.fullpath + postfix, detail=False):
         meta = load_meta(
-            filepath, fs=loc.fs, load_value=False, follow_links=False
+            filepath,
+            fs=loc.fs,
+            load_value=False,
+            follow_links=False,
+            try_migrations=try_migrations,
         )
         type_ = meta.__class__
         if isinstance(meta, MlemLink):

mlem/core/objects.py

@@ -150,6 +150,7 @@ def read(
         cls: Type[T],
         location: Location,
         follow_links: bool = True,
+        try_migrations: bool = False,
     ) -> T:
         """
         Read object in (path, fs)
@@ -170,6 +171,10 @@ def read(
         )
         with location.open() as f:
             payload = safe_load(f)
+        if try_migrations:
+            from mlem.api.migrations import apply_migrations
+
+            payload, _ = apply_migrations(payload)
         res = parse_obj_as(MlemObject, payload).bind(location)
         if follow_links and isinstance(res, MlemLink):
             link = res.load_link()

tests/api/test_migrations.py

@@ -61,3 +61,14 @@ def test_directory(tmpdir, old_data, new_data, recursive):
             assert load_meta(subdir_path) != new_data
         except ValidationError:
             pass
+
+
+@pytest.mark.parametrize("old_data,new_data", [model_03])
+def test_load_with_migration(tmpdir, old_data, new_data):
+    path = tmpdir / "model.mlem"
+    path.write_text(safe_dump(old_data), encoding="utf8")
+
+    meta = load_meta(path, try_migrations=True)
+
+    assert isinstance(meta, MlemObject)
+    assert meta == new_data