Fetch all entries when listing datasets (#62)

* fetch all entries in list method and improve get by ID/Name methods

* but back release condition

Author: nathanthorpe

.github/workflows/package.yml

@@ -40,6 +40,7 @@ jobs:
           poetry build
 
       - name: Publish python package
+        if: github.event_name == 'release'
         run: |
           poetry publish
         env:

cirro/api/services/base.py

@@ -1,9 +1,34 @@
 from abc import ABC
+from typing import Dict
 
 from cirro.api.clients import ApiClient
 from cirro.api.config import AppConfig
 
 
+def fetch_all_items(client: ApiClient, query: str, input_variables: Dict, batch_size=10000, max_items=None):
+    """
+    Fetches all items from a paginated graphql api
+    """
+    next_token = None
+    items = []
+    while True:
+        variables = {
+            'nextToken': next_token,
+            'limit': batch_size,
+            **input_variables
+        }
+        resp = client.query(query, variables)
+        query_name = next(iter(resp.keys()))
+        items.extend(resp[query_name]['items'])
+
+        next_token = resp[query_name]['nextToken']
+        if not next_token:
+            return items
+
+        if max_items and len(items) >= max_items:
+            return items
+
+
 class BaseService(ABC):
     _api_client: ApiClient
     _configuration: AppConfig

cirro/api/services/dataset.py

@@ -1,12 +1,13 @@
 import json
 import logging
 import uuid
-from typing import List, Union
+from typing import List, Union, Optional
 
 from cirro.api.clients.utils import filter_deleted
 from cirro.api.models.dataset import CreateIngestDatasetInput, DatasetCreateResponse, Dataset
 from cirro.api.models.file import FileAccessContext, File
 from cirro.api.models.status import Status
+from cirro.api.services.base import fetch_all_items
 from cirro.api.services.file import FileEnabledService
 
 logger = logging.getLogger()
@@ -56,17 +57,40 @@ def find_by_project(self, project_id: str, name: str = None) -> List[Dataset]:
                 'status': {
                     'eq': Status.COMPLETED.value
                 }
-            },
-            # TODO: Implement pagination
-            'limit': 10000
+            }
         }
         if name:
             variables['filter']['name'] = {'eq': name}
 
-        resp = self._api_client.query(query, variables=variables)['datasetsByProject']
-        not_deleted = filter_deleted(resp['items'])
+        items = fetch_all_items(self._api_client, query, variables)
+        not_deleted = filter_deleted(items)
         return [Dataset.from_record(item) for item in not_deleted]
 
+    def get_from_id(self, _id: str) -> Optional[Dataset]:
+        query = '''
+             query GetDataset($id: ID!) {
+                getDataset(id: $id) {
+                  id
+                  status
+                  name
+                  desc
+                  sourceDatasets
+                  paramJson
+                  infoJson
+                  process
+                  project
+                  createdAt
+                  updatedAt
+                }
+              }
+              '''
+
+        item = self._api_client.query(query, variables={'id': _id})['getDataset']
+        if not item:
+            return
+
+        return Dataset.from_record(item)
+
     def create(self, create_request: CreateIngestDatasetInput) -> DatasetCreateResponse:
         """
         Creates an ingested dataset.

cirro/sdk/project.py

@@ -51,12 +51,18 @@ def list_datasets(self) -> DataPortalDatasets:
     def get_dataset_by_name(self, name: str = None) -> DataPortalDataset:
         """Return the dataset with the specified name."""
 
-        return self.list_datasets().get_by_name(name)
+        dataset = next(iter(self._client.dataset.find_by_project(self.id, name=name)), None)
+        if dataset is None:
+            raise DataPortalAssetNotFound(f'Dataset with name {name} not found')
+        return DataPortalDataset(dataset, self._client)
 
     def get_dataset_by_id(self, _id: str = None) -> DataPortalDataset:
         """Return the dataset with the specified id."""
 
-        return self.list_datasets().get_by_id(_id)
+        dataset = self._client.dataset.get_from_id(_id=_id)
+        if dataset is None:
+            raise DataPortalAssetNotFound(f'Dataset with ID {_id} not found')
+        return DataPortalDataset(dataset, self._client)
 
     def list_references(self, reference_type: str = None) -> DataPortalReferences:
         """