Add query jobs (#189)

* Add list jobs

* Use QueryJobs

* address comments

* handle (str, str) case better

Author: corentinmusard

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.36.0] - 2025-05-27
+
+## Added
+
+- Added `query` method to `JobClient` to query jobs in a given temporal extent and filter by automation id
+
 ## [0.35.0] - 2025-04-29
 
 ## Added
@@ -161,7 +167,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Released packages: `tilebox-datasets`, `tilebox-workflows`, `tilebox-storage`, `tilebox-grpc`
 
 
-[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.35.0...HEAD
+[Unreleased]: https://github.com/tilebox/tilebox-python/compare/v0.36.0...HEAD
+[0.36.0]: https://github.com/tilebox/tilebox-python/compare/v0.35.0...v0.36.0
 [0.35.0]: https://github.com/tilebox/tilebox-python/compare/v0.34.0...v0.35.0
 [0.34.0]: https://github.com/tilebox/tilebox-python/compare/v0.33.1...v0.34.0
 [0.33.1]: https://github.com/tilebox/tilebox-python/compare/v0.33.0...v0.33.1

tilebox-datasets/tests/data/datapoint.py

@@ -35,6 +35,7 @@
     AnyMessage,
     Datapoint,
     DatapointInterval,
+    DatapointIntervalLike,
     DatapointPage,
     IngestResponse,
     QueryResultPage,
@@ -59,6 +60,19 @@ def datapoint_intervals(draw: DrawFn) -> DatapointInterval:
     return DatapointInterval(start, end, start_exclusive, end_inclusive)
 
 
+@composite
+def datapoint_intervals_like(draw: DrawFn) -> DatapointIntervalLike:
+    """A hypothesis strategy for generating random datapoint intervals"""
+    interval = draw(datapoint_intervals())
+    return draw(
+        one_of(
+            just(interval),
+            just((str(interval.start_id), str(interval.end_id))),
+            just((interval.start_id, interval.end_id)),
+        )
+    )
+
+
 @composite
 def example_datapoints(draw: DrawFn, generated_fields: bool = False, missing_fields: bool = False) -> ExampleDatapoint:
     """

tilebox-datasets/tests/data/test_datapoint.py

@@ -3,6 +3,7 @@
 from tests.data.datapoint import (
     anys,
     datapoint_intervals,
+    datapoint_intervals_like,
     datapoint_pages,
     datapoints,
     ingest_datapoints_responses,
@@ -13,6 +14,7 @@
     AnyMessage,
     Datapoint,
     DatapointInterval,
+    DatapointIntervalLike,
     DatapointPage,
     IngestResponse,
     QueryResultPage,
@@ -25,6 +27,19 @@ def test_datapoint_intervals_to_message_and_back(interval: DatapointInterval) ->
     assert DatapointInterval.from_message(interval.to_message()) == interval
 
 
+@given(datapoint_intervals_like())
+def test_parse_datapoint_interval_from_tuple(interval: DatapointIntervalLike) -> None:
+    parsed = DatapointInterval.parse(interval)
+
+    if isinstance(interval, DatapointInterval):
+        assert parsed == interval, f"Failed parsing interval from {interval}"
+        assert parsed.start_exclusive == interval.start_exclusive
+        assert parsed.end_inclusive == interval.end_inclusive
+    else:
+        assert not parsed.start_exclusive
+        assert parsed.end_inclusive
+
+
 @given(anys())
 def test_anys_to_message_and_back(any_: AnyMessage) -> None:
     assert AnyMessage.from_message(any_.to_message()) == any_

tilebox-datasets/tilebox/datasets/aio/dataset.py

@@ -7,17 +7,18 @@
 import xarray as xr
 from tqdm.auto import tqdm
 
+from _tilebox.grpc.aio.pagination import Pagination as PaginationProtocol
+from _tilebox.grpc.aio.pagination import paginated_request
 from _tilebox.grpc.aio.producer_consumer import async_producer_consumer
 from _tilebox.grpc.error import ArgumentError, NotFoundError
 from tilebox.datasets.aio.pagination import (
-    paginated_request,
     with_progressbar,
     with_time_progress_callback,
     with_time_progressbar,
 )
 from tilebox.datasets.data.collection import CollectionInfo
 from tilebox.datasets.data.data_access import QueryFilters, SpatialFilter, SpatialFilterLike
-from tilebox.datasets.data.datapoint import DatapointInterval, DatapointPage, QueryResultPage
+from tilebox.datasets.data.datapoint import DatapointInterval, DatapointIntervalLike, DatapointPage, QueryResultPage
 from tilebox.datasets.data.datasets import Dataset
 from tilebox.datasets.data.pagination import Pagination
 from tilebox.datasets.data.time_interval import TimeInterval, TimeIntervalLike
@@ -242,7 +243,7 @@ async def _find_legacy(self, datapoint_id: str, skip_data: bool = False) -> xr.D
 
     async def _find_interval(
         self,
-        datapoint_id_interval: tuple[str, str] | tuple[UUID, UUID],
+        datapoint_id_interval: DatapointIntervalLike,
         end_inclusive: bool = True,
         *,
         skip_data: bool = False,
@@ -266,18 +267,13 @@ async def _find_interval(
                 datapoint_id_interval, end_inclusive, skip_data=skip_data, show_progress=show_progress
             )
 
-        start_id, end_id = datapoint_id_interval
-
         filters = QueryFilters(
-            temporal_extent=DatapointInterval(
-                start_id=as_uuid(start_id),
-                end_id=as_uuid(end_id),
-                start_exclusive=False,
-                end_inclusive=end_inclusive,
-            )
+            temporal_extent=DatapointInterval.parse(datapoint_id_interval, end_inclusive=end_inclusive)
         )
 
-        request = partial(self._dataset._service.query, [self._collection.id], filters, skip_data)
+        async def request(page: PaginationProtocol) -> QueryResultPage:
+            query_page = Pagination(page.limit, page.starting_after)
+            return await self._dataset._service.query([self._collection.id], filters, skip_data, query_page)
 
         initial_page = Pagination()
         pages = paginated_request(request, initial_page)
@@ -288,27 +284,19 @@ async def _find_interval(
 
     async def _find_interval_legacy(
         self,
-        datapoint_id_interval: tuple[str, str] | tuple[UUID, UUID],
+        datapoint_id_interval: DatapointIntervalLike,
         end_inclusive: bool = True,
         *,
         skip_data: bool = False,
         show_progress: bool = False,
     ) -> xr.Dataset:
-        start_id, end_id = datapoint_id_interval
+        datapoint_interval = DatapointInterval.parse(datapoint_id_interval, end_inclusive=end_inclusive)
 
-        datapoint_interval = DatapointInterval(
-            start_id=as_uuid(start_id),
-            end_id=as_uuid(end_id),
-            start_exclusive=False,
-            end_inclusive=end_inclusive,
-        )
-        request = partial(
-            self._dataset._service.get_dataset_for_datapoint_interval,
-            str(self._collection.id),
-            datapoint_interval,
-            skip_data,
-            False,
-        )
+        async def request(page: PaginationProtocol) -> DatapointPage:
+            query_page = Pagination(page.limit, page.starting_after)
+            return await self._dataset._service.get_dataset_for_datapoint_interval(
+                str(self._collection.id), datapoint_interval, skip_data, False, query_page
+            )
 
         initial_page = Pagination()
         pages = paginated_request(request, initial_page)
@@ -427,9 +415,10 @@ async def _iter_pages(
             yield page
 
     async def _load_page(
-        self, filters: QueryFilters, skip_data: bool, page: Pagination | None = None
+        self, filters: QueryFilters, skip_data: bool, page: PaginationProtocol | None = None
     ) -> QueryResultPage:
-        return await self._dataset._service.query([self._collection.id], filters, skip_data, page)
+        query_page = Pagination(page.limit, page.starting_after) if page else Pagination()
+        return await self._dataset._service.query([self._collection.id], filters, skip_data, query_page)
 
     async def _load_legacy(
         self,
@@ -472,10 +461,11 @@ async def _iter_pages_legacy(
             yield page
 
     async def _load_page_legacy(
-        self, time_interval: TimeInterval, skip_data: bool, skip_meta: bool, page: Pagination | None = None
+        self, time_interval: TimeInterval, skip_data: bool, skip_meta: bool, page: PaginationProtocol | None = None
     ) -> DatapointPage:
+        query_page = Pagination(page.limit, page.starting_after) if page else Pagination()
         return await self._dataset._service.get_dataset_for_time_interval(
-            str(self._collection.id), time_interval, skip_data, skip_meta, page
+            str(self._collection.id), time_interval, skip_data, skip_meta, query_page
         )
 
     async def ingest(

tilebox-datasets/tilebox/datasets/aio/pagination.py

@@ -1,5 +1,5 @@
 import time
-from collections.abc import AsyncIterator, Awaitable, Callable
+from collections.abc import AsyncIterator
 from datetime import datetime, timezone
 from typing import TypeVar
 
@@ -9,41 +9,11 @@
     TimeInterval,
 )
 from tilebox.datasets.data.datapoint import DatapointPage, QueryResultPage
-from tilebox.datasets.data.pagination import Pagination
 from tilebox.datasets.progress import ProgressCallback, TimeIntervalProgressBar
 
 ResultPage = TypeVar("ResultPage", bound=DatapointPage | QueryResultPage)
 
 
-async def paginated_request(
-    paging_request: Callable[[Pagination], Awaitable[ResultPage]],
-    initial_page: Pagination | None = None,
-) -> AsyncIterator[ResultPage]:
-    """Make a paginated request to a gRPC service endpoint.
-
-    The endpoint is expected to return a next_page field, which is used for subsequent requests. Once no such
-    next_page field is returned, the request is completed.
-
-    Args:
-        paging_request: A function that takes a page as input and returns a Datapoints object
-            Often this will be a functools.partial object that wraps a gRPC service endpoint
-            and only leaves the page argument remaining
-        initial_page: The initial page to request
-
-    Yields:
-        Datapoints: The individual pages of the response
-    """
-    if initial_page is None:
-        initial_page = Pagination()
-
-    response = await paging_request(initial_page)
-    yield response
-
-    while response.next_page.starting_after is not None:
-        response = await paging_request(response.next_page)
-        yield response
-
-
 async def with_progressbar(
     paginated_request: AsyncIterator[ResultPage],
     progress_description: str,

tilebox-datasets/tilebox/datasets/data/datapoint.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any
+from typing import Any, TypeAlias
 from uuid import UUID
 
 from tilebox.datasets.data.pagination import Pagination
@@ -9,6 +9,8 @@
 from tilebox.datasets.datasetsv1 import core_pb2, data_access_pb2, data_ingestion_pb2
 from tilebox.datasets.message_pool import get_message_type
 
+DatapointIntervalLike: TypeAlias = "tuple[str, str] | tuple[UUID, UUID] | DatapointInterval"
+
 
 @dataclass(frozen=True)
 class DatapointInterval:
@@ -34,6 +36,47 @@ def to_message(self) -> core_pb2.DatapointInterval:
             end_inclusive=self.end_inclusive,
         )
 
+    @classmethod
+    def parse(
+        cls, arg: DatapointIntervalLike, start_exclusive: bool = False, end_inclusive: bool = True
+    ) -> "DatapointInterval":
+        """
+        Convert a variety of input types to a DatapointInterval.
+
+        Supported input types:
+        - DatapointInterval: Return the input as is
+        - tuple of two UUIDs: Return an DatapointInterval with start and end id set to the given values
+        - tuple of two strings: Return an DatapointInterval with start and end id set to the UUIDs parsed from the given strings
+
+        Args:
+            arg: The input to convert
+            start_exclusive: Whether the start id is exclusive
+            end_inclusive: Whether the end id is inclusive
+
+        Returns:
+            DatapointInterval: The parsed ID interval
+        """
+
+        match arg:
+            case DatapointInterval(_, _, _, _):
+                return arg
+            case (UUID(), UUID()):
+                start, end = arg
+                return DatapointInterval(
+                    start_id=start,
+                    end_id=end,
+                    start_exclusive=start_exclusive,
+                    end_inclusive=end_inclusive,
+                )
+            case (str(), str()):
+                start, end = arg
+                return DatapointInterval(
+                    start_id=UUID(start),
+                    end_id=UUID(end),
+                    start_exclusive=start_exclusive,
+                    end_inclusive=end_inclusive,
+                )
+
 
 @dataclass(frozen=True)
 class AnyMessage: