fix: get_dataset and list_dataset to consider source workspace slug (HEXA-1366)

openhexa/sdk/datasets/dataset.py

@@ -370,11 +370,13 @@ def __init__(
         slug: str,
         name: str,
         description: str,
+        source_workspace_slug: str = None,
     ):
         self.id = id
         self.slug = slug
         self.name = name
         self.description = description
+        self.source_workspace_slug = source_workspace_slug
 
     def create_version(self, name: typing.Any) -> DatasetVersion:
         """Build a dataset version, save it and return it."""
@@ -457,7 +459,7 @@ def versions(self) -> VersionsIterator:
 
     def __repr__(self) -> str:
         """Safe representation of the dataset."""
-        return f"<Dataset slug={self.slug} id={self.id}>"
+        return f"<Dataset slug={self.slug} id={self.id} source_workspace_slug={self.source_workspace_slug}>"
 
 
 class FileNotFound(Exception):

openhexa/sdk/workspaces/current_workspace.py

@@ -521,14 +521,17 @@ def create_dataset(self, name: str, description: str):
         identifier = rsp["createDataset"]["dataset"]["slug"]
         return self.get_dataset(identifier)
 
-    def get_dataset(self, identifier: str) -> Dataset:
+    def get_dataset(self, identifier: str, source_workspace_slug: str = None) -> Dataset:
         """Get a dataset by its identifier.
 
         Parameters
         ----------
         identifier : str
             The identifier of the dataset in the OpenHEXA backend
 
+        source_workspace_slug : str
+            The slug of the workspace the dataset belongs to,  defaults to the current workspace slug
+
         Returns
         -------
         Dataset
@@ -544,10 +547,6 @@ def get_dataset(self, identifier: str) -> Dataset:
             query getDataset($datasetSlug: String!, $workspaceSlug: String!) {
                 datasetLinkBySlug(datasetSlug: $datasetSlug, workspaceSlug: $workspaceSlug) {
                     id
-                    workspace {
-                        slug
-                        name
-                    }
                     dataset {
                         id
                         slug
@@ -558,22 +557,33 @@ def get_dataset(self, identifier: str) -> Dataset:
                             name
                             description
                         }
+                        workspace {
+                            slug
+                        }
                     }
                 }
             }
         """,
-            {"datasetSlug": identifier, "workspaceSlug": self.slug},
+            {"datasetSlug": identifier, "workspaceSlug": source_workspace_slug or self.slug},
         )
         data = response["datasetLinkBySlug"]
 
         if data is None:
-            raise ValueError(f"Dataset {identifier} does not exist.")
+            raise ValueError(
+                f"Dataset {identifier} does not exist on workspace {source_workspace_slug or self.slug}."
+                + (
+                    " If you try and get a dataset shared from another workspace, please provide the source_workspace_slug parameter."
+                    if not source_workspace_slug
+                    else ""
+                )
+            )
 
         return Dataset(
             id=data["dataset"]["id"],
             slug=data["dataset"]["slug"],
             name=data["dataset"]["name"],
             description=data["dataset"]["description"],
+            source_workspace_slug=data["dataset"]["workspace"]["slug"],
         )
 
     def list_datasets(self) -> list[Dataset]:
@@ -595,6 +605,9 @@ def list_datasets(self) -> list[Dataset]:
                                 slug
                                 name
                                 description
+                                workspace {
+                                    slug
+                                }
                             }
                         }
 
@@ -611,6 +624,7 @@ def list_datasets(self) -> list[Dataset]:
                 slug=d["dataset"]["slug"],
                 name=d["dataset"]["name"],
                 description=d["dataset"]["description"],
+                source_workspace_slug=d["dataset"]["workspace"]["slug"],
             )
             for d in data
         ]

tests/test_dataset.py

@@ -52,13 +52,12 @@ def graphql_responses(url, request):
                         "data": {
                             "datasetLinkBySlug": {
                                 "id": "link_id",
-                                "workspace": "WORKSPACE",
                                 "dataset": {
                                     "id": "dataset_id",
                                     "slug": "my-dataset",
                                     "name": "My Dataset",
                                     "description": "My Dataset description",
-                                    "createdAt": "2021-01-01T00:00:00.000Z",
+                                    "workspace": {"slug": "source-workspace"},
                                 },
                             }
                         }