From ade2404a79a336720c50be13e43d396eeeef9938 Mon Sep 17 00:00:00 2001 From: omaraldulimay <120497558+omaraldulimay@users.noreply.github.com> Date: Fri, 13 Nov 2026 15:00:29 +0100 Subject: [PATCH 1/2] STACAPIJobDatabase: faster counting, paging and existence check (#972) Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 2 + openeo/extra/job_management/_manager.py | 6 +- openeo/extra/job_management/stac_job_db.py | 31 ++++++++- .../extra/job_management/test_stac_job_db.py | 64 +++++++++++++++++-- 4 files changed, 94 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 743f5bf29..dce3ee4d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- `STACAPIJobDatabase`: improve performance by using `numberMatched` counts, a larger search page size and a direct collection lookup ([#972](https://github.com/Open-EO/openeo-python-client/issues/972)) + ### Removed ### Fixed diff --git a/openeo/extra/job_management/_manager.py b/openeo/extra/job_management/_manager.py index c1f059240..5590816aa 100644 --- a/openeo/extra/job_management/_manager.py +++ b/openeo/extra/job_management/_manager.py @@ -384,7 +384,8 @@ def run_loop(): stats["run_jobs loop"] += 1 # Show current stats and sleep - _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}") + if _log.isEnabledFor(logging.INFO): + _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}") for _ in range(int(max(1, self.poll_sleep))): time.sleep(1) if self._stop_thread: @@ -526,7 +527,8 @@ def run_jobs( stats["run_jobs loop"] += 1 # Show current stats and sleep - _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}") + if _log.isEnabledFor(logging.INFO): + _log.info(f"Job status histogram: {job_db.count_by_status()}. Run stats: {dict(stats)}") time.sleep(self.poll_sleep) stats["sleep"] += 1 diff --git a/openeo/extra/job_management/stac_job_db.py b/openeo/extra/job_management/stac_job_db.py index f6f7d0e8c..6af298671 100644 --- a/openeo/extra/job_management/stac_job_db.py +++ b/openeo/extra/job_management/stac_job_db.py @@ -3,6 +3,7 @@ import concurrent.futures import datetime import logging +import warnings from typing import Iterable, List, Optional, Union import geopandas as gpd @@ -58,9 +59,12 @@ def __init__( self.geometry_column = geometry_column self.base_url = stac_root_url self.bulk_size = 500 + # Page size to request from the STAC API when searching items (servers may cap this to a lower maximum). + self.search_page_size = 1000 def exists(self) -> bool: - return any(c.id == self.collection_id for c in self.client.get_collections()) + # Look up the collection directly instead of paging through all collections of the STAC API. + return self.client.get_collection(self.collection_id) is not None def _normalize_df(self, df: pd.DataFrame) -> pd.DataFrame: """ @@ -167,12 +171,32 @@ def count_by_status(self, statuses: Iterable[str] = ()) -> dict: if isinstance(statuses, str): statuses = {statuses} statuses = set(statuses) + if statuses: + # Ask the STAC API for the number of matching items, instead of downloading all of them. + counts = {s: self._count_matched(status_filter=self._status_filter([s])) for s in statuses} + if all(c is not None for c in counts.values()): + return counts items = self.get_by_status(statuses) if items is None: return {k: 0 for k in statuses} else: return items["status"].value_counts().to_dict() + @staticmethod + def _status_filter(statuses: Iterable[str]) -> Optional[str]: + return " OR ".join([f"\"properties.status\"='{s}'" for s in statuses]) or None + + def _count_matched(self, status_filter: Optional[str]) -> Optional[int]: + """ + Get the number of matching items as reported by the STAC API (``numberMatched`` or ``context.matched``), + or ``None`` if the STAC API does not report it. + """ + search = self.client.search(method="GET", collections=[self.collection_id], filter=status_filter, limit=1) + with warnings.catch_warnings(): + # pystac_client warns when the count is not supported, we handle that by falling back. + warnings.simplefilter("ignore") + return search.matched() + def _search_result_to_df(self, search_result: pystac_client.ItemSearch) -> pd.DataFrame: """Build a DataFrame from a STAC ItemSearch result.""" series = [self.series_from(item) for item in search_result.items()] @@ -184,12 +208,12 @@ def get_by_status(self, statuses: Iterable[str], max: Optional[int] = None) -> p statuses = {statuses} statuses = set(statuses) - status_filter = " OR ".join([f"\"properties.status\"='{s}'" for s in statuses]) if statuses else None search_results = self.client.search( method="GET", collections=[self.collection_id], - filter=status_filter, + filter=self._status_filter(statuses), max_items=max, + limit=self.search_page_size, ) df = self._search_result_to_df(search_results) @@ -204,6 +228,7 @@ def get_by_indices(self, indices: Iterable[Union[int, str]]) -> pd.DataFrame: method="GET", collections=[self.collection_id], ids=[str(i) for i in indices], + limit=self.search_page_size, ) df = self._search_result_to_df(search_results) return df diff --git a/tests/extra/job_management/test_stac_job_db.py b/tests/extra/job_management/test_stac_job_db.py index 274e197b6..ea72450d3 100644 --- a/tests/extra/job_management/test_stac_job_db.py +++ b/tests/extra/job_management/test_stac_job_db.py @@ -27,6 +27,9 @@ def mock_pystac_client(): MagicMock(id="collection-1"), MagicMock(id="collection-2"), ] + mock_client.get_collection.side_effect = lambda collection_id: ( + MagicMock(id=collection_id) if collection_id in {"collection-1", "collection-2"} else None + ) with patch("pystac_client.Client.open", return_value=mock_client): yield mock_client @@ -340,17 +343,56 @@ def test_count_by_status(self, mock_get_by_status, job_db_exists): ) assert job_db_exists.count_by_status() == {"not_started": 1} + def test_exists_does_not_list_all_collections(self, job_db_exists): + assert job_db_exists.exists() is True + job_db_exists.client.get_collection.assert_called_once_with("collection-1") + job_db_exists.client.get_collections.assert_not_called() + + def test_count_by_status_uses_number_matched(self, requests_mock): + dummy_stac_api = DummyStacApi(root_url="http://stacapi.test", requests_mock=requests_mock) + for i, status in enumerate(["not_started", "running", "not_started", "finished"]): + dummy_stac_api.predefine_item( + collection_id="collection-123", + item=_pystac_item(id=f"item-{i}", properties={"status": status}), + ) + job_db = STACAPIJobDatabase(collection_id="collection-123", stac_root_url="http://stacapi.test") + + counts = job_db.count_by_status(["not_started", "running", "queued"]) + + assert counts == {"not_started": 2, "running": 1, "queued": 0} + searches = [r for r in requests_mock.request_history if r.path == "/search"] + # One search per status, each only asking for a single item: no full download of the job database. + assert len(searches) == 3 + assert all(r.qs["limit"] == ["1"] for r in searches) + + def test_count_by_status_falls_back_without_number_matched(self, requests_mock): + dummy_stac_api = DummyStacApi( + root_url="http://stacapi.test", requests_mock=requests_mock, support_number_matched=False + ) + for i, status in enumerate(["not_started", "running", "not_started"]): + dummy_stac_api.predefine_item( + collection_id="collection-123", + item=_pystac_item(id=f"item-{i}", properties={"status": status}), + ) + job_db = STACAPIJobDatabase(collection_id="collection-123", stac_root_url="http://stacapi.test") + + assert job_db.count_by_status(["not_started", "running"]) == {"not_started": 2, "running": 1} + def test_get_by_status_no_filter(self, job_db_exists): job_db_exists.get_by_status(()) job_db_exists.client.search.assert_called_once_with( - method="GET", collections=["collection-1"], filter=None, max_items=None + method="GET", collections=["collection-1"], filter=None, max_items=None, limit=1000 ) def test_get_by_status_with_filter(self, job_db_exists): job_db_exists.get_by_status(["not_started"]) job_db_exists.client.search.assert_called_once_with( - method="GET", collections=["collection-1"], filter="\"properties.status\"='not_started'", max_items=None + method="GET", + collections=["collection-1"], + filter="\"properties.status\"='not_started'", + max_items=None, + limit=1000, ) def test_get_by_status_result(self, requests_mock): @@ -491,14 +533,16 @@ def sleep_mock(): class DummyStacApi: """Minimal dummy implementation of a STAC API for testing purposes.""" - def __init__(self, *, root_url: str = "http://stacapi.test", requests_mock): + def __init__(self, *, root_url: str = "http://stacapi.test", requests_mock, support_number_matched: bool = True): self.root_url = root_url.rstrip("/") self._requests_mock = requests_mock + self._support_number_matched = support_number_matched requests_mock.get(f"{self.root_url}/", json=self._get_root()) self.collections: List[dict] = [] requests_mock.get(f"{self.root_url}/collections", json=self._get_collections) requests_mock.post(f"{self.root_url}/collections", json=self._post_collections) + requests_mock.get(re.compile(rf"{re.escape(self.root_url)}/collections/[^/]+$"), json=self._get_collection) self.items: Dict[str, Dict[str, Any]] = collections.defaultdict(dict) requests_mock.post( @@ -527,6 +571,15 @@ def _get_collections(self, request, context): """Handler of `GET /collections` requests.""" return {"collections": self.collections} + def _get_collection(self, request, context): + """Handler of `GET /collections/{collection_id}` requests.""" + collection_id = request.path.split("/")[-1] + for collection in self.collections: + if collection["id"] == collection_id: + return collection + context.status_code = 404 + return {"code": "NotFound", "description": f"Collection {collection_id} not found"} + def _post_collections(self, request, context): """Handler of `POST /collections` requests.""" post_data = request.json() @@ -581,11 +634,14 @@ def _get_search(self, request, context): statuses = set(re.findall(r"\"properties\.status\"='(\w+)'", property_filter)) items = [i for i in items if i.get("properties", {}).get("status") in statuses] - return { + result = { "type": "FeatureCollection", "features": items, "links": [], } + if self._support_number_matched: + result["numberMatched"] = len(items) + return result def test_run_jobs_basic(tmp_path, dummy_backend_foo, requests_mock, sleep_mock): From 34999c984da1597986e71b4badac71f48f922f01 Mon Sep 17 00:00:00 2001 From: omaraldulimay <120497558+omaraldulimay@users.noreply.github.com> Date: Sat, 14 Nov 2026 14:18:14 +0100 Subject: [PATCH 2/2] STACAPIJobDatabase.exists(): handle 404 error of older pystac_client versions Co-Authored-By: Claude Opus 5.5 --- openeo/extra/job_management/stac_job_db.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/openeo/extra/job_management/stac_job_db.py b/openeo/extra/job_management/stac_job_db.py index 6af298671..6ab5d3f31 100644 --- a/openeo/extra/job_management/stac_job_db.py +++ b/openeo/extra/job_management/stac_job_db.py @@ -18,6 +18,7 @@ try: # Optional dependecy import pystac_client + import pystac_client.exceptions except ImportError: pass @@ -64,7 +65,13 @@ def __init__( def exists(self) -> bool: # Look up the collection directly instead of paging through all collections of the STAC API. - return self.client.get_collection(self.collection_id) is not None + try: + return self.client.get_collection(self.collection_id) is not None + except pystac_client.exceptions.APIError as e: + # Older pystac_client versions raise on a missing collection instead of returning None. + if getattr(e, "status_code", None) == 404: + return False + raise def _normalize_df(self, df: pd.DataFrame) -> pd.DataFrame: """