Send the datasource when downloading a blob

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:

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"
 

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}"