diff --git a/dagshub/data_engine/model/datasource.py b/dagshub/data_engine/model/datasource.py index bbeab214..18a38e15 100644 --- a/dagshub/data_engine/model/datasource.py +++ b/dagshub/data_engine/model/datasource.py @@ -337,6 +337,11 @@ def all(self, load_documents=True, load_annotations=True) -> "QueryResult": ds = self.limit(None) return ds.fetch(load_documents=load_documents, load_annotations=load_annotations) + def get_datapoint(self, path: str) -> Datapoint: + """Get a datapoint by its path.""" + result = (self[self["path"] == path]).head(1) + return result[path] + def select(self, *selected: Union[str, Field]) -> "Datasource": """ Select which fields should appear in the query result. diff --git a/tests/data_engine/test_datasource.py b/tests/data_engine/test_datasource.py index e6f6e0dc..d9a20a65 100644 --- a/tests/data_engine/test_datasource.py +++ b/tests/data_engine/test_datasource.py @@ -280,3 +280,49 @@ def test_annotation_in_dataframe(ds): def test_annotation_in_dataframe_to_new_datasource(ds, other_ds): _test_dataframe_annotation_addition(ds, other_ds) + + +def test_get_datapoint_by_path(ds): + datapoint_path = "folder/file.jpg" + expected = Datapoint( + datasource=ds, + path=datapoint_path, + datapoint_id=123, + metadata={}, + ) + query_result = QueryResult( + _entries=[expected], + datasource=ds, + fields=[], + ) + ds.source.client.head.return_value = query_result + + actual = ds.get_datapoint(datapoint_path) + + assert actual is expected + + queried_ds, size = ds.source.client.head.call_args.args + assert size == 1 + assert queried_ds.get_query().filter.tree_to_dict() == { + "eq": { + "data": { + "field": "path", + "value": datapoint_path, + } + } + } + + +def test_get_datapoint_by_path_not_found(ds): + datapoint_path = "missing/file.jpg" + query_result = QueryResult( + _entries=[], + datasource=ds, + fields=[], + ) + ds.source.client.head.return_value = query_result + + with pytest.raises(KeyError) as exc_info: + ds.get_datapoint(datapoint_path) + + assert exc_info.value.args == (datapoint_path,)