From b243535352fae3a8b95a2e7ebb9d54adbf0f5bd7 Mon Sep 17 00:00:00 2001 From: Guy Smoilovsky Date: Fri, 17 Jul 2026 00:15:45 +0300 Subject: [PATCH 1/4] Add datasource copy API --- dagshub/data_engine/client/data_client.py | 6 +++ dagshub/data_engine/client/gql_mutations.py | 18 ++++++++ dagshub/data_engine/datasources.py | 26 ++++++++++++ dagshub/data_engine/model/datasource.py | 10 +++++ tests/data_engine/test_datasource_copy.py | 46 +++++++++++++++++++++ 5 files changed, 106 insertions(+) create mode 100644 tests/data_engine/test_datasource_copy.py diff --git a/dagshub/data_engine/client/data_client.py b/dagshub/data_engine/client/data_client.py index fa8bacb17..9f182421d 100644 --- a/dagshub/data_engine/client/data_client.py +++ b/dagshub/data_engine/client/data_client.py @@ -83,6 +83,12 @@ def create_datasource(self, ds: "DatasourceState") -> DatasourceResult: res = self._exec(q, params) return dacite.from_dict(DatasourceResult, res["createDatasource"], config=dacite_config) + def copy_datasource(self, source: Union[int, str], name: str, as_of: Optional[int] = None) -> DatasourceResult: + q = GqlMutations.copy_datasource() + params = GqlMutations.copy_datasource_params(source=source, name=name, as_of=as_of) + res = self._exec(q, params) + return dacite.from_dict(DatasourceResult, res["copyDatasource"], config=dacite_config) + def head(self, datasource: "Datasource", size: Optional[int] = None) -> QueryResult: """ Retrieve a subset of data from the datasource headers. diff --git a/dagshub/data_engine/client/gql_mutations.py b/dagshub/data_engine/client/gql_mutations.py index 1b53f9fde..6372dd9c3 100644 --- a/dagshub/data_engine/client/gql_mutations.py +++ b/dagshub/data_engine/client/gql_mutations.py @@ -27,6 +27,24 @@ def create_datasource(): def create_datasource_params(name: str, url: str, ds_type: DatasourceType): return {"name": name, "url": url, "dsType": str(ds_type.value)} + @staticmethod + @functools.lru_cache() + def copy_datasource(): + return ( + GqlQuery() + .operation( + "mutation", + name="copyDatasource", + input={"$source": "ID!", "$name": "String!", "$asOf": "DateTime"}, + ) + .query("copyDatasource", input={"source": "$source", "name": "$name", "asOf": "$asOf"}) + .fields(["id", "name", "rootUrl", "integrationStatus", "preprocessingStatus", "type"]) + ) + + @staticmethod + def copy_datasource_params(source: Union[int, str], name: str, as_of: Optional[int]): + return {"source": source, "name": name, "asOf": as_of} + @staticmethod @functools.lru_cache() def update_metadata(): diff --git a/dagshub/data_engine/datasources.py b/dagshub/data_engine/datasources.py index e2c3be75b..10154fc41 100644 --- a/dagshub/data_engine/datasources.py +++ b/dagshub/data_engine/datasources.py @@ -1,3 +1,4 @@ +import datetime import json import logging from typing import Optional, Union, List, TYPE_CHECKING, Dict @@ -5,6 +6,7 @@ from dagshub.common.analytics import send_analytics_event from dagshub.common.api.repo import RepoAPI from dagshub.common.util import lazy_load, removeprefix +from dagshub.common.util import to_timestamp from dagshub.data_engine.client.data_client import DataClient from dagshub.data_engine.model.datasource import Datasource from dagshub.data_engine.model.datasource_state import DatasourceState, DatasourceType, path_regexes @@ -60,6 +62,29 @@ def create(*args, **kwargs) -> Datasource: return create_datasource(*args, **kwargs) +def copy_datasource( + repo: str, + source: Union[str, int, Datasource], + name: str, + as_of: Optional[Union[float, int, datetime.datetime]] = None, +) -> Datasource: + """Create a new datasource from the source's latest state or a point-in-time snapshot. + + The copy runs asynchronously. Its preprocessing status exposes progress through the + same APIs used for newly scanned datasources. Version history before the selected + snapshot is not copied. + """ + source_id = source.source.id if isinstance(source, Datasource) else source + if source_id is None: + raise ValueError("source datasource must have an id") + result = DataClient(repo).copy_datasource( + source=source_id, + name=name, + as_of=to_timestamp(as_of) if as_of is not None else None, + ) + return Datasource(DatasourceState.from_gql_result(repo, result)) + + def get_or_create(repo: str, name: str, path: str, revision: Optional[str] = None) -> Datasource: """ First attempts to get the repo datasource with the given name, and only if that fails, @@ -236,6 +261,7 @@ def _load_datasources_from_run( __all__ = [ create_datasource.__name__, create.__name__, + copy_datasource.__name__, create_from_bucket.__name__, create_from_repo.__name__, get_datasource.__name__, diff --git a/dagshub/data_engine/model/datasource.py b/dagshub/data_engine/model/datasource.py index bbeab214e..fa1b38582 100644 --- a/dagshub/data_engine/model/datasource.py +++ b/dagshub/data_engine/model/datasource.py @@ -173,6 +173,16 @@ def __init__( self.ngrok_listener = None + def copy( + self, + name: str, + as_of: Optional[Union[float, int, datetime.datetime]] = None, + ) -> "Datasource": + """Create an asynchronous copy of this datasource at the requested snapshot.""" + from dagshub.data_engine.datasources import copy_datasource + + return copy_datasource(self.source.repo, self, name, as_of) + @property def has_explicit_context(self): return self._explicit_update_ctx is not None diff --git a/tests/data_engine/test_datasource_copy.py b/tests/data_engine/test_datasource_copy.py new file mode 100644 index 000000000..9e2a49d14 --- /dev/null +++ b/tests/data_engine/test_datasource_copy.py @@ -0,0 +1,46 @@ +from dagshub.data_engine.client.gql_mutations import GqlMutations +from dagshub.data_engine.client.data_client import DataClient +from dagshub.data_engine.client.models import PreprocessingStatus + + +def test_copy_datasource_mutation_and_params(): + query = GqlMutations.copy_datasource().generate() + + assert "mutation copyDatasource" in query + assert "$source: ID!" in query + assert "$asOf: DateTime" in query + assert "copyDatasource(source: $source, name: $name, asOf: $asOf)" in query + assert GqlMutations.copy_datasource_params(12, "snapshot", 1_700_000_000) == { + "source": 12, + "name": "snapshot", + "asOf": 1_700_000_000, + } + + +def test_copy_datasource_params_accept_latest_state(): + assert GqlMutations.copy_datasource_params("12", "latest", None)["asOf"] is None + + +def test_data_client_returns_copied_datasource(monkeypatch): + client = object.__new__(DataClient) + captured = {} + + def fake_exec(query, params): + captured.update(params) + return { + "copyDatasource": { + "id": "13", + "name": "snapshot", + "rootUrl": "repo://owner/repo/main:data", + "integrationStatus": "VALID", + "preprocessingStatus": "IN_PROGRESS", + "type": "REPOSITORY", + } + } + + monkeypatch.setattr(client, "_exec", fake_exec) + result = client.copy_datasource(12, "snapshot", 1_700_000_000) + + assert captured == {"source": 12, "name": "snapshot", "asOf": 1_700_000_000} + assert result.id == "13" + assert result.preprocessingStatus == PreprocessingStatus.IN_PROGRESS From bcd93cad2892638b830167437a1af93ddd021d4a Mon Sep 17 00:00:00 2001 From: Guy Smoilovsky Date: Fri, 17 Jul 2026 09:59:42 +0300 Subject: [PATCH 2/4] Materialize datasource queries --- dagshub/data_engine/client/data_client.py | 6 ++- dagshub/data_engine/client/gql_mutations.py | 22 ++++++-- dagshub/data_engine/client/models.py | 14 +++++ dagshub/data_engine/datasources.py | 18 ++++--- dagshub/data_engine/model/datasource.py | 5 +- dagshub/data_engine/model/datasource_state.py | 10 +++- tests/data_engine/test_datasource_copy.py | 52 ++++++++++++++++--- 7 files changed, 103 insertions(+), 24 deletions(-) diff --git a/dagshub/data_engine/client/data_client.py b/dagshub/data_engine/client/data_client.py index 9f182421d..32fec4183 100644 --- a/dagshub/data_engine/client/data_client.py +++ b/dagshub/data_engine/client/data_client.py @@ -83,9 +83,11 @@ def create_datasource(self, ds: "DatasourceState") -> DatasourceResult: res = self._exec(q, params) return dacite.from_dict(DatasourceResult, res["createDatasource"], config=dacite_config) - def copy_datasource(self, source: Union[int, str], name: str, as_of: Optional[int] = None) -> DatasourceResult: + def copy_datasource( + self, source: Union[int, str], name: str, query_input: Optional[Dict[str, Any]] = None + ) -> DatasourceResult: q = GqlMutations.copy_datasource() - params = GqlMutations.copy_datasource_params(source=source, name=name, as_of=as_of) + params = GqlMutations.copy_datasource_params(source=source, name=name, query=query_input) res = self._exec(q, params) return dacite.from_dict(DatasourceResult, res["copyDatasource"], config=dacite_config) diff --git a/dagshub/data_engine/client/gql_mutations.py b/dagshub/data_engine/client/gql_mutations.py index 6372dd9c3..6f784777b 100644 --- a/dagshub/data_engine/client/gql_mutations.py +++ b/dagshub/data_engine/client/gql_mutations.py @@ -35,15 +35,27 @@ def copy_datasource(): .operation( "mutation", name="copyDatasource", - input={"$source": "ID!", "$name": "String!", "$asOf": "DateTime"}, + input={"$source": "ID!", "$name": "String!", "$query": "QueryInput"}, ) - .query("copyDatasource", input={"source": "$source", "name": "$name", "asOf": "$asOf"}) - .fields(["id", "name", "rootUrl", "integrationStatus", "preprocessingStatus", "type"]) + .query("copyDatasource", input={"source": "$source", "name": "$name", "query": "$query"}) + .fields( + [ + "id", + "name", + "rootUrl", + "integrationStatus", + "preprocessingStatus", + "type", + "origin {sourceDatasourceId sourceRepoId sourceName sourceRootUrl sourceType " + "sourceBackingType creatorId createdAt query}", + ] + ) + .param_validator(Validators.query_input_validator("query")) ) @staticmethod - def copy_datasource_params(source: Union[int, str], name: str, as_of: Optional[int]): - return {"source": source, "name": name, "asOf": as_of} + def copy_datasource_params(source: Union[int, str], name: str, query: Optional[Dict[str, Any]]): + return {"source": source, "name": name, "query": query} @staticmethod @functools.lru_cache() diff --git a/dagshub/data_engine/client/models.py b/dagshub/data_engine/client/models.py index 9e0f24271..4cee9cfb8 100644 --- a/dagshub/data_engine/client/models.py +++ b/dagshub/data_engine/client/models.py @@ -101,6 +101,19 @@ def from_metadata_field_schema(mfs: MetadataFieldSchema) -> "MetadataSelectField ) +@dataclass +class DatasourceOriginResult: + sourceDatasourceId: Union[str, int] + sourceRepoId: Union[str, int] + sourceName: str + sourceRootUrl: str + sourceType: DatasourceType + sourceBackingType: str + creatorId: Union[str, int] + createdAt: int + query: str + + @dataclass class DatasourceResult: id: Union[str, int] @@ -110,6 +123,7 @@ class DatasourceResult: preprocessingStatus: PreprocessingStatus type: DatasourceType metadataFields: Optional[List[MetadataFieldSchema]] + origin: Optional[DatasourceOriginResult] = None @dataclass diff --git a/dagshub/data_engine/datasources.py b/dagshub/data_engine/datasources.py index 10154fc41..76d0312ca 100644 --- a/dagshub/data_engine/datasources.py +++ b/dagshub/data_engine/datasources.py @@ -66,21 +66,27 @@ def copy_datasource( repo: str, source: Union[str, int, Datasource], name: str, - as_of: Optional[Union[float, int, datetime.datetime]] = None, + query: Optional[Dict] = None, ) -> Datasource: - """Create a new datasource from the source's latest state or a point-in-time snapshot. + """Materialize a new datasource from a query over an existing datasource. The copy runs asynchronously. Its preprocessing status exposes progress through the - same APIs used for newly scanned datasources. Version history before the selected - snapshot is not copied. + same APIs used for newly scanned datasources. When ``source`` is a Datasource, its + current filters, projection, limit, and as-of timestamp are used automatically. + Version history before the materialized state is not copied. """ - source_id = source.source.id if isinstance(source, Datasource) else source + if isinstance(source, Datasource): + source_id = source.source.id + if query is None: + query = source.serialize_gql_query_input() + else: + source_id = source if source_id is None: raise ValueError("source datasource must have an id") result = DataClient(repo).copy_datasource( source=source_id, name=name, - as_of=to_timestamp(as_of) if as_of is not None else None, + query_input=query, ) return Datasource(DatasourceState.from_gql_result(repo, result)) diff --git a/dagshub/data_engine/model/datasource.py b/dagshub/data_engine/model/datasource.py index fa1b38582..24cd5e606 100644 --- a/dagshub/data_engine/model/datasource.py +++ b/dagshub/data_engine/model/datasource.py @@ -176,12 +176,11 @@ def __init__( def copy( self, name: str, - as_of: Optional[Union[float, int, datetime.datetime]] = None, ) -> "Datasource": - """Create an asynchronous copy of this datasource at the requested snapshot.""" + """Asynchronously materialize this datasource's current query as a datasource.""" from dagshub.data_engine.datasources import copy_datasource - return copy_datasource(self.source.repo, self, name, as_of) + return copy_datasource(self.source.repo, self, name) @property def has_explicit_context(self): diff --git a/dagshub/data_engine/model/datasource_state.py b/dagshub/data_engine/model/datasource_state.py index 93f50e081..efd5105c1 100644 --- a/dagshub/data_engine/model/datasource_state.py +++ b/dagshub/data_engine/model/datasource_state.py @@ -7,7 +7,13 @@ from os import PathLike from dagshub.common.api.repo import RepoAPI, PathNotFoundError from dagshub.data_engine.client.data_client import DataClient -from dagshub.data_engine.client.models import DatasourceType, DatasourceResult, PreprocessingStatus, MetadataFieldSchema +from dagshub.data_engine.client.models import ( + DatasourceType, + DatasourceResult, + DatasourceOriginResult, + PreprocessingStatus, + MetadataFieldSchema, +) from dagshub.data_engine.model.datapoint import Datapoint from dagshub.data_engine.model.errors import DatasourceAlreadyExistsError, DatasourceNotFoundError from dagshub.common.util import multi_urljoin @@ -43,6 +49,7 @@ class DatasourceState: client: DataClient = field(init=False) repoApi: RepoAPI = field(init=False) metadata_fields: List[MetadataFieldSchema] = field(init=False) + origin: Optional[DatasourceOriginResult] = field(init=False, default=None) _revision: Optional[str] = field(init=False, default=None) @@ -209,6 +216,7 @@ def _update_from_ds_result(self, ds: DatasourceResult): self.source_type = ds.type self.preprocessing_status = ds.preprocessingStatus self.metadata_fields = [] if ds.metadataFields is None else ds.metadataFields + self.origin = ds.origin if self.source_type == DatasourceType.REPOSITORY: self.revision = self.path_parts()["revision"] diff --git a/tests/data_engine/test_datasource_copy.py b/tests/data_engine/test_datasource_copy.py index 9e2a49d14..e2a654dd7 100644 --- a/tests/data_engine/test_datasource_copy.py +++ b/tests/data_engine/test_datasource_copy.py @@ -1,6 +1,8 @@ from dagshub.data_engine.client.gql_mutations import GqlMutations from dagshub.data_engine.client.data_client import DataClient from dagshub.data_engine.client.models import PreprocessingStatus +from dagshub.data_engine.model.datasource import Datasource +from dagshub.data_engine.model.datasource_state import DatasourceState def test_copy_datasource_mutation_and_params(): @@ -8,17 +10,18 @@ def test_copy_datasource_mutation_and_params(): assert "mutation copyDatasource" in query assert "$source: ID!" in query - assert "$asOf: DateTime" in query - assert "copyDatasource(source: $source, name: $name, asOf: $asOf)" in query - assert GqlMutations.copy_datasource_params(12, "snapshot", 1_700_000_000) == { + assert "$query: QueryInput" in query + assert "copyDatasource(source: $source, name: $name, query: $query)" in query + query_input = {"select": [{"name": "score", "alias": "confidence"}], "limit": 20} + assert GqlMutations.copy_datasource_params(12, "snapshot", query_input) == { "source": 12, "name": "snapshot", - "asOf": 1_700_000_000, + "query": query_input, } def test_copy_datasource_params_accept_latest_state(): - assert GqlMutations.copy_datasource_params("12", "latest", None)["asOf"] is None + assert GqlMutations.copy_datasource_params("12", "latest", None)["query"] is None def test_data_client_returns_copied_datasource(monkeypatch): @@ -35,12 +38,47 @@ def fake_exec(query, params): "integrationStatus": "VALID", "preprocessingStatus": "IN_PROGRESS", "type": "REPOSITORY", + "origin": { + "sourceDatasourceId": "12", + "sourceRepoId": "7", + "sourceName": "source", + "sourceRootUrl": "repo://owner/repo/main:data", + "sourceType": "REPOSITORY", + "sourceBackingType": "postgres", + "creatorId": "3", + "createdAt": 1_700_000_001, + "query": '{"asOf":1700000000,"limit":20}', + }, } } monkeypatch.setattr(client, "_exec", fake_exec) - result = client.copy_datasource(12, "snapshot", 1_700_000_000) + query_input = {"asOf": 1_700_000_000, "limit": 20} + result = client.copy_datasource(12, "snapshot", query_input) - assert captured == {"source": 12, "name": "snapshot", "asOf": 1_700_000_000} + assert captured == {"source": 12, "name": "snapshot", "query": query_input} assert result.id == "13" assert result.preprocessingStatus == PreprocessingStatus.IN_PROGRESS + assert result.origin is not None + assert result.origin.sourceName == "source" + assert result.origin.sourceBackingType == "postgres" + + +def test_datasource_copy_materializes_current_query(monkeypatch): + state = object.__new__(DatasourceState) + state.repo = "owner/repo" + state.id = 12 + queried = Datasource(state).select("score").limit(5) + captured = {} + + def fake_copy(repo, source, name): + captured.update(repo=repo, source=source, name=name, query=source.serialize_gql_query_input()) + return "copied" + + monkeypatch.setattr("dagshub.data_engine.datasources.copy_datasource", fake_copy) + + assert queried.copy("projection") == "copied" + assert captured["repo"] == "owner/repo" + assert captured["source"] is queried + assert captured["name"] == "projection" + assert captured["query"] == {"select": [{"name": "score"}], "query": None, "limit": 5} From de3fd90069b4f602ed5f39950580ba4268605709 Mon Sep 17 00:00:00 2001 From: Guy Smoilovsky Date: Fri, 17 Jul 2026 10:09:12 +0300 Subject: [PATCH 3/4] Remove obsolete copy imports --- dagshub/data_engine/datasources.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/dagshub/data_engine/datasources.py b/dagshub/data_engine/datasources.py index 76d0312ca..918f739ed 100644 --- a/dagshub/data_engine/datasources.py +++ b/dagshub/data_engine/datasources.py @@ -1,4 +1,3 @@ -import datetime import json import logging from typing import Optional, Union, List, TYPE_CHECKING, Dict @@ -6,7 +5,6 @@ from dagshub.common.analytics import send_analytics_event from dagshub.common.api.repo import RepoAPI from dagshub.common.util import lazy_load, removeprefix -from dagshub.common.util import to_timestamp from dagshub.data_engine.client.data_client import DataClient from dagshub.data_engine.model.datasource import Datasource from dagshub.data_engine.model.datasource_state import DatasourceState, DatasourceType, path_regexes From 09c6881fd0b1add09302f50312a5c4eb28749f14 Mon Sep 17 00:00:00 2001 From: Guy Smoilovsky Date: Fri, 17 Jul 2026 11:22:53 +0300 Subject: [PATCH 4/4] Rename datasource copy method --- dagshub/data_engine/model/datasource.py | 2 +- tests/data_engine/test_datasource_copy.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/dagshub/data_engine/model/datasource.py b/dagshub/data_engine/model/datasource.py index 24cd5e606..2a0fd51fc 100644 --- a/dagshub/data_engine/model/datasource.py +++ b/dagshub/data_engine/model/datasource.py @@ -173,7 +173,7 @@ def __init__( self.ngrok_listener = None - def copy( + def copy_to( self, name: str, ) -> "Datasource": diff --git a/tests/data_engine/test_datasource_copy.py b/tests/data_engine/test_datasource_copy.py index e2a654dd7..6393c293c 100644 --- a/tests/data_engine/test_datasource_copy.py +++ b/tests/data_engine/test_datasource_copy.py @@ -77,7 +77,7 @@ def fake_copy(repo, source, name): monkeypatch.setattr("dagshub.data_engine.datasources.copy_datasource", fake_copy) - assert queried.copy("projection") == "copied" + assert queried.copy_to("projection") == "copied" assert captured["repo"] == "owner/repo" assert captured["source"] is queried assert captured["name"] == "projection"