diff --git a/dagshub/data_engine/model/datasource_state.py b/dagshub/data_engine/model/datasource_state.py index 93f50e08..ff280995 100644 --- a/dagshub/data_engine/model/datasource_state.py +++ b/dagshub/data_engine/model/datasource_state.py @@ -121,7 +121,13 @@ def blob_path(self, sha: str) -> str: """ Returns the path for the blob of a datasource """ - return f"{self.repoApi.data_engine_url}/blob/{sha}" + url = f"{self.repoApi.data_engine_url}/blob/{sha}" + if self.id is None: + return url + # Blobs are stored per datasource. The server can find the blob without this, but only + # by searching the repo's datasources, so tell it where to look. Older servers ignore + # the unknown parameter and still serve the blob. + return f"{url}?datasource={self.id}" @cached_property def root_content_path(self) -> str: diff --git a/tests/data_engine/annotation_import/test_annotation_parsing.py b/tests/data_engine/annotation_import/test_annotation_parsing.py index 133311c0..c3a67759 100644 --- a/tests/data_engine/annotation_import/test_annotation_parsing.py +++ b/tests/data_engine/annotation_import/test_annotation_parsing.py @@ -3,6 +3,7 @@ from pathlib import Path from typing import Union from unittest.mock import MagicMock +from urllib.parse import urlsplit import pytest from dagshub_annotation_converter.ir.image import IRSegmentationImageAnnotation @@ -60,7 +61,9 @@ def mock_annotation_query_result( def mock_get_blob(*args, **kwargs) -> Union[bytes, PathLike]: download_url: str = args[0] - blob_hash = download_url.split("/")[-1] + # The hash is the last path segment; the url may also carry a ?datasource= hint, which + # the real server reads separately and which is not part of the hash. + blob_hash = urlsplit(download_url).path.split("/")[-1] load_into_memory = args[4] blob_path = _res_folder / f"{blob_hash}.json" diff --git a/tests/data_engine/test_datasource_state.py b/tests/data_engine/test_datasource_state.py index 9920cbc1..a4af2cbf 100644 --- a/tests/data_engine/test_datasource_state.py +++ b/tests/data_engine/test_datasource_state.py @@ -84,3 +84,19 @@ def test_bucket_regex_incorrect(in_str, mock_dagshub_auth): ds.source_type = DatasourceType.REPOSITORY with pytest.raises(InvalidPathFormatError): ds.path_parts() + + +def test_blob_path_includes_datasource(ds): + source = ds.source + sha = "a" * 64 + + assert source.blob_path(sha) == f"{source.repoApi.data_engine_url}/blob/{sha}?datasource={source.id}" + + +def test_blob_path_without_datasource_id(mock_dagshub_auth): + # The server falls back to searching the repo's datasources when the hint is absent, + # so an unsaved datasource must still produce a usable url rather than "datasource=None". + source = DatasourceState(repo="user/repo") + sha = "a" * 64 + + assert source.blob_path(sha) == f"{source.repoApi.data_engine_url}/blob/{sha}"