From e7860db407913da9e01c195bec13d22cf98086e3 Mon Sep 17 00:00:00 2001 From: Guy Smoilovsky Date: Wed, 15 Jul 2026 17:29:52 +0300 Subject: [PATCH] Point blob downloads at their datasource Blobs are now stored under a per-datasource prefix on the server. The server can still find a blob without knowing its datasource, but only by searching the repo's datasources one by one, so send the datasource we already have. This is purely an optimization: servers predating the parameter ignore it, and a datasource with no id yet (never saved) omits it and is served by the search path. Also stop the annotation test's blob mock from treating the query string as part of the blob hash. It derives the hash from the download url, unlike the real client, which caches by hash independently of the url. --- dagshub/data_engine/model/datasource_state.py | 8 +++++++- .../annotation_import/test_annotation_parsing.py | 5 ++++- tests/data_engine/test_datasource_state.py | 16 ++++++++++++++++ 3 files changed, 27 insertions(+), 2 deletions(-) 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}"