fix: gaps in versioned objects for dataset objects

packages/syft-datasets/src/syft_datasets/dataset_manager.py

@@ -2,7 +2,6 @@
 
 from typing_extensions import Self
 
-import yaml
 from syft_migration import ProtocolSchema
 
 from .types import PathLike, to_path
@@ -362,5 +361,4 @@ def _private_config_without_data_dir(self, ref: DatasetRef) -> bytes:
         """
         config = self.storage.read_private_config(ref)
         config.data_dir = Path("")
-        data = config.disk_dict()
-        return yaml.safe_dump(data, indent=2, sort_keys=False).encode()
+        return self.storage.wire_bytes(ref, config)

packages/syft-datasets/src/syft_datasets/dataset_storage.py

@@ -529,11 +529,25 @@ def write_private_config(
         codec = self._codec_for(ref.protocol_version)
         return self._write(codec, codec.private_metadata_path(ref), config, ref)
 
+    def wire_bytes(self, ref: DatasetRef, obj: MigratableObject) -> bytes:
+        """The bytes to put on the wire for ``obj`` under this ref's protocol.
+
+        Same shaping as the write methods: downgrade to the protocol schema,
+        then the codec's on-disk format. Use it when the bytes go to a peer
+        instead of to a path.
+        """
+        codec = self._codec_for(ref.protocol_version)
+        return codec.dumps(self._downgrade(obj, ref)).encode()
+
     # -- internals -----------------------------------------------------------
     def _upgrade(self, data: dict, canonical_name: str) -> MigratableObject:
         obj = self.service.load(data)
         return self.service.migrate(obj, self.registry.latest_version(canonical_name))
 
+    def _downgrade(self, obj: MigratableObject, ref: DatasetRef) -> MigratableObject:
+        schema = self.registry.schema_for_protocol_version(ref.protocol_version)
+        return self.service.migrate_to_schema(obj, schema)
+
     def _write(
         self,
         codec: ProtocolCodec,
@@ -542,7 +556,5 @@ def _write(
         ref: DatasetRef,
     ) -> Path:
         """Downgrade ``obj`` to the ref's protocol schema, then let the codec persist it."""
-        schema = self.registry.schema_for_protocol_version(ref.protocol_version)
-        downgraded = self.service.migrate_to_schema(obj, schema)
-        codec.write(path, downgraded)
+        codec.write(path, self._downgrade(obj, ref))
         return path

packages/syft-datasets/src/syft_datasets/models/dataset/v1.py

@@ -6,7 +6,6 @@
 from typing import ClassVar
 from uuid import UUID, uuid4
 
-import yaml
 from pydantic import Field, PrivateAttr
 from syft_migration import MigratableObject
 from syft_notebook_ui.formatter_mixin import (
@@ -23,7 +22,7 @@
 from ...dataset_ref import DatasetRef
 from ...migrations import dataset_registry
 from ...url import SyftBoxURL
-from ..private_dataset_config.v1 import PrivateDatasetConfigV1
+from ..private_dataset_config import PrivateDatasetConfig
 
 
 def _utcnow() -> datetime:
@@ -99,25 +98,30 @@ def get_readme(self) -> str | None:
     def mock_dir(self) -> Path:
         return self._url_to_path(self.mock_url)
 
-    @property
-    def private_config_path(self) -> Path:
+    def _require_owner(self, what: str) -> None:
         if self.syftbox_config.email != self.owner:
             raise ValueError(
-                "Cannot access private config for a dataset owned by another user."
+                f"Cannot access {what} for a dataset owned by another user."
             )
+
+    @property
+    def private_config_path(self) -> Path:
+        self._require_owner("private config")
         return self._private_metadata_dir / PRIVATE_METADATA_FILENAME
 
     @cached_property
-    def private_config(self) -> PrivateDatasetConfigV1:
-        config_path = self.private_config_path
-        if not config_path.exists():
-            raise FileNotFoundError(
-                f"Private dataset config not found at {config_path}"
-            )
-        data = yaml.safe_load(config_path.read_text()) or {}
-        data.setdefault("canonical_name", "PrivateDatasetConfig")
-        data.setdefault("version", "1")
-        return PrivateDatasetConfigV1(**data)
+    def private_config(self) -> PrivateDatasetConfig:
+        """This dataset's private config, upgraded to the latest version.
+
+        Raises ValueError if another user owns the dataset, and
+        PrivateConfigNotFoundError if the file is absent.
+        """
+        # Local import: dataset_storage imports this module.
+        from ...dataset_storage import DatasetStorage
+
+        self._require_owner("private config")
+        storage = DatasetStorage(config=self.syftbox_config)
+        return storage.read_private_config(self._ref)
 
     @property
     def private_dir(self) -> Path:
@@ -136,10 +140,7 @@ def private_dir(self) -> Path:
 
     @property
     def _private_metadata_dir(self) -> Path:
-        if self.syftbox_config.email != self.owner:
-            raise ValueError(
-                "Cannot access private data for a dataset owned by another user."
-            )
+        self._require_owner("private data")
         return self.private_dir
 
     @property

packages/syft-datasets/src/syft_datasets/protocolcodecs/base.py

@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import Iterator
 
+import yaml
 from syft_migration import MigratableObject
 
 from ..config import SyftBoxConfig
@@ -51,4 +52,20 @@ def iter_dataset_refs(self, datasite_email: str) -> Iterator[DatasetRef]:
     def read(self, path: Path, canonical_name: str) -> dict: ...
 
     @abstractmethod
-    def write(self, path: Path, obj: MigratableObject) -> None: ...
+    def _data_for_disk(self, obj: MigratableObject) -> dict:
+        """The fields to serialize for this layout.
+
+        The one place a codec decides its on-disk shape. Abstract on purpose: a
+        codec states its shaping rather than inheriting a default that may not
+        match the layout it speaks.
+        """
+        ...
+
+    def dumps(self, obj: MigratableObject) -> str:
+        """YAML for ``obj`` in this codec's on-disk format."""
+        return yaml.safe_dump(self._data_for_disk(obj), indent=2, sort_keys=False)
+
+    def write(self, path: Path, obj: MigratableObject) -> None:
+        """Persist ``obj`` at ``path`` in this codec's on-disk format."""
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(self.dumps(obj))

packages/syft-datasets/src/syft_datasets/protocolcodecs/v0.py

@@ -74,10 +74,9 @@ def read(self, path: Path, canonical_name: str) -> dict:
         data.setdefault("version", "1")
         return data
 
-    def write(self, path: Path, obj: MigratableObject) -> None:
+    def _data_for_disk(self, obj: MigratableObject) -> dict:
         data = obj.disk_dict()
         # Byte-match the pre-versioning (<= 0.1.20) on-disk format.
         data.pop("canonical_name", None)
         data.pop("version", None)
-        path.parent.mkdir(parents=True, exist_ok=True)
-        path.write_text(yaml.safe_dump(data, indent=2, sort_keys=False))
+        return data

packages/syft-datasets/src/syft_datasets/protocolcodecs/v1.py

@@ -98,7 +98,6 @@ def read(self, path: Path, canonical_name: str) -> dict:
         # Files already carry canonical_name/version on disk.
         return yaml.safe_load(path.read_text()) or {}
 
-    def write(self, path: Path, obj: MigratableObject) -> None:
-        data = obj.disk_dict()
-        path.parent.mkdir(parents=True, exist_ok=True)
-        path.write_text(yaml.safe_dump(data, indent=2, sort_keys=False))
+    def _data_for_disk(self, obj: MigratableObject) -> dict:
+        # This layout keeps the identity fields on disk.
+        return obj.disk_dict()

packages/syft-datasets/tests/migrations/p2p/test_current_protocol.py

@@ -170,6 +170,40 @@ def test_multi_version_write_for_mixed_audience(tmp_path: Path):
     assert mgr.get("demo")._ref.protocol_version == "1"
 
 
+def _wire_private_metadata(mgr: SyftDatasetManager, name: str, protocol_version: str):
+    files = mgr.get_private_dataset_files(name, protocol_version=protocol_version)
+    metas = [
+        content
+        for path, content in files.items()
+        if path.name == "private_metadata.yaml"
+    ]
+    assert len(metas) == 1
+    return yaml.safe_load(metas[0])
+
+
+def test_private_files_wire_private_metadata_in_protocol_format(tmp_path: Path):
+    """get_private_dataset_files serializes private_metadata through the codec."""
+    mock, private, readme = _create_dataset_files(tmp_path)
+    mgr = _dataset_manager(tmp_path)
+
+    mgr.create(name="flat", mock_path=mock, private_path=private, readme_path=readme)
+    raw0 = _wire_private_metadata(mgr, "flat", "0")
+    assert "canonical_name" not in raw0 and "version" not in raw0
+    assert "uid" in raw0
+    assert not Path(str(raw0["data_dir"])).is_absolute()
+
+    mgr.create(
+        name="nested",
+        mock_path=mock,
+        private_path=private,
+        readme_path=readme,
+        protocol_versions=["1"],
+    )
+    raw1 = _wire_private_metadata(mgr, "nested", "1")
+    assert raw1["canonical_name"] == "PrivateDatasetConfig" and raw1["version"] == "1"
+    assert not Path(str(raw1["data_dir"])).is_absolute()
+
+
 def test_delete_removes_all_protocol_versions(tmp_path: Path):
     schema0 = dataset_registry.schema_for_protocol_version("0")
     schema1 = dataset_registry.schema_for_protocol_version("1")

packages/syft-datasets/tests/migrations/p2p/test_protocol_codecs.py

@@ -18,7 +18,7 @@
 from syft_datasets.dataset_storage import DatasetRef, DatasetStorage
 from syft_datasets.migrations import dataset_registry
 from syft_datasets.migrations.registry import DATASET_PROTOCOL_VERSION
-from syft_datasets.models import Dataset
+from syft_datasets.models import Dataset, PrivateDatasetConfig
 from syft_datasets.protocolcodecs import ProtocolCodec
 from syft_datasets.protocolcodecs.v1 import DatasetConfigV1
 from syft_datasets.url import SyftBoxURL
@@ -123,6 +123,92 @@ def test_v0_writes_flat_no_identity_v1_nests_with_identity(tmp_path: Path):
         assert loaded.version == dataset_registry.latest_version("Dataset")
 
 
+def test_wire_bytes_match_write_and_strip_private_config_identity_on_v0(tmp_path: Path):
+    storage = _storage(tmp_path)
+    ref0 = DatasetRef(DO_EMAIL, "flat", "0")
+    ref1 = DatasetRef(DO_EMAIL, "nested", "1")
+    dataset = _mock_dataset(storage, ref0)
+    config = PrivateDatasetConfig(uid=dataset.uid, data_dir=Path("/secret/local/path"))
+
+    p0 = storage.write_private_config(ref0, config)
+    p1 = storage.write_private_config(ref1, config)
+
+    assert storage.wire_bytes(ref0, config) == p0.read_bytes()
+    assert storage.wire_bytes(ref1, config) == p1.read_bytes()
+
+    raw0 = yaml.safe_load(p0.read_text())
+    assert "canonical_name" not in raw0 and "version" not in raw0
+    assert list(raw0) == ["uid", "data_dir"]
+
+    raw1 = yaml.safe_load(p1.read_text())
+    assert raw1["canonical_name"] == "PrivateDatasetConfig" and raw1["version"] == "1"
+
+
+@pytest.fixture
+def private_config_v2():
+    """A throwaway PrivateDatasetConfig V2, as the next release would ship it.
+
+    Subclassing a registered class registers it, so this lands in the global
+    dataset_registry; the teardown pops it back out, because the registry has no
+    deregister call and another test must not see a version "2".
+    """
+    from syft_datasets.models import PrivateDatasetConfigV1
+
+    # The class body registers V2, so everything after it needs the teardown.
+    class PrivateDatasetConfigV2(PrivateDatasetConfigV1, registry=dataset_registry):
+        version: str = "2"
+        checksum: str = ""
+
+    try:
+        dataset_registry.register_migration(
+            canonical_name="PrivateDatasetConfig",
+            from_version="1",
+            to_version="2",
+            fn=lambda obj: PrivateDatasetConfigV2(
+                **obj.model_dump(exclude={"version"})
+            ),
+        )
+        yield PrivateDatasetConfigV2
+    finally:
+        dataset_registry.objects["PrivateDatasetConfig"].pop("2", None)
+        dataset_registry.migrations.get("PrivateDatasetConfig", {}).pop(
+            ("1", "2"), None
+        )
+
+
+def test_dataset_private_config_upgrades_a_v1_file_to_the_latest_version(
+    tmp_path: Path, private_config_v2
+):
+    # Dataset.private_config must read through storage, which upgrades. Reading the
+    # file into PrivateDatasetConfigV1 by name would return version 1 forever.
+    storage = _storage(tmp_path)
+    ref = DatasetRef(DO_EMAIL, "demo", "1")
+    dataset = _mock_dataset(storage, ref)
+    storage.write_dataset_metadata(ref, dataset)
+
+    path = storage.private_metadata_path(ref)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(
+        yaml.safe_dump(
+            {
+                "canonical_name": "PrivateDatasetConfig",
+                "version": "1",
+                "uid": str(dataset.uid),
+                "data_dir": "",
+            },
+            indent=2,
+            sort_keys=False,
+        )
+    )
+
+    config = storage.read_dataset(ref).private_config
+
+    assert isinstance(config, private_config_v2)
+    assert config.version == "2"
+    # The uid comes from the file, so a fresh V2 that ignored it would fail here.
+    assert config.uid == dataset.uid
+
+
 # -- behavior 2: each codec scans only its own layout --------------------------
 def test_scan_partitions_by_layout(tmp_path: Path):
     storage = _storage(tmp_path)