Control canonical links in load_stac_from_job

CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Allow `Connection.load_stac_from_job()` to require or avoid canonical result links. ([#634](https://github.com/Open-EO/openeo-python-client/issues/634))
+
 ### Changed
 
 ### Removed

openeo/rest/connection.py

@@ -19,6 +19,7 @@
     Iterable,
     Iterator,
     List,
+    Literal,
     Mapping,
     Optional,
     Sequence,
@@ -1478,48 +1479,64 @@ def load_stac_from_job(
         temporal_extent: Union[Sequence[InputDate], Parameter, str, None] = None,
         bands: Optional[List[str]] = None,
         properties: Optional[Dict[str, Union[str, PGNode, Callable]]] = None,
+        *,
+        use_canonical_link: Literal["auto", "require", "avoid"] = "auto",
     ) -> DataCube:
         """
         Convenience function to directly load the results of a finished openEO job
         (as a STAC collection) with :py:meth:`load_stac` in a new openEO process graph.
 
-        When available, the "canonical" link (signed URL) of the job results will be used.
+        By default, the "canonical" link (signed URL) of the job results is used when available.
 
         :param job: a :py:class:`~openeo.rest.job.BatchJob` or job id pointing to a finished job.
             Note that the  :py:class:`~openeo.rest.job.BatchJob` approach allows to point
             to a batch job on a different back-end.
         :param spatial_extent: limit data to specified bounding box or polygons
         :param temporal_extent: limit data to specified temporal interval.
         :param bands: limit data to the specified bands
+        :param use_canonical_link: control use of the canonical result link:
+            ``"auto"`` prefers it but falls back to the regular job results URL,
+            ``"require"`` fails if it is unavailable, and ``"avoid"`` always uses
+            the regular job results URL.
 
         .. versionadded:: 0.30.0
+
+        .. versionchanged:: 0.52.0
+            Added the ``use_canonical_link`` argument.
         """
-        # TODO #634 add option to require or avoid the canonical link
+        if use_canonical_link not in {"auto", "require", "avoid"}:
+            raise ValueError(f"Invalid use_canonical_link mode {use_canonical_link!r}")
+
         if isinstance(job, str):
             job = BatchJob(job_id=job, connection=self)
         elif not isinstance(job, BatchJob):
             raise ValueError("job must be a BatchJob or job id")
 
-        try:
-            job_results = job.get_results()
-
-            canonical_links = [
-                link["href"]
-                for link in job_results.get_metadata().get("links", [])
-                if link.get("rel") == "canonical" and "href" in link
-            ]
-            if len(canonical_links) == 0:
-                _log.warning("No canonical link found in job results metadata. Using job results URL instead.")
-                stac_link = job.get_results_metadata_url(full=True)
-            else:
-                if len(canonical_links) > 1:
-                    _log.warning(
-                        f"Multiple canonical links found in job results metadata: {canonical_links}. Picking first one."
-                    )
-                stac_link = canonical_links[0]
-        except OpenEoApiError as e:
-            _log.warning(f"Failed to get the canonical job results: {e!r}. Using job results URL instead.")
-            stac_link = job.get_results_metadata_url(full=True)
+        stac_link = job.get_results_metadata_url(full=True)
+        if use_canonical_link != "avoid":
+            try:
+                job_results = job.get_results()
+
+                canonical_links = [
+                    link["href"]
+                    for link in job_results.get_metadata().get("links", [])
+                    if link.get("rel") == "canonical" and "href" in link
+                ]
+                if len(canonical_links) == 0:
+                    message = "No canonical link found in job results metadata."
+                    if use_canonical_link == "require":
+                        raise OpenEoClientException(message)
+                    _log.warning(f"{message} Using job results URL instead.")
+                else:
+                    if len(canonical_links) > 1:
+                        _log.warning(
+                            f"Multiple canonical links found in job results metadata: {canonical_links}. Picking first one."
+                        )
+                    stac_link = canonical_links[0]
+            except OpenEoApiError as e:
+                if use_canonical_link == "require":
+                    raise
+                _log.warning(f"Failed to get the canonical job results: {e!r}. Using job results URL instead.")
 
         return self.load_stac(
             url=stac_link,

tests/rest/test_connection.py

@@ -3241,7 +3241,8 @@ def test_property_filtering(self, con120, properties, expected):
             }
         }
 
-    def test_load_stac_from_job_canonical(self, con120, requests_mock):
+    @pytest.mark.parametrize("use_canonical_link", ["auto", "require"])
+    def test_load_stac_from_job_canonical(self, con120, requests_mock, use_canonical_link):
         requests_mock.get(
             API_URL + "jobs/j0bi6/results",
             json={
@@ -3258,7 +3259,7 @@ def test_load_stac_from_job_canonical(self, con120, requests_mock):
             },
         )
         job = con120.job("j0bi6")
-        cube = con120.load_stac_from_job(job)
+        cube = con120.load_stac_from_job(job, use_canonical_link=use_canonical_link)
 
         fg = cube.flat_graph()
         assert fg == {
@@ -3295,6 +3296,30 @@ def test_load_stac_from_job_unsigned(self, con120, requests_mock):
             }
         }
 
+    def test_load_stac_from_job_require_canonical_missing(self, con120, requests_mock):
+        requests_mock.get(
+            API_URL + "jobs/j0bi6/results",
+            json={"links": [{"href": "https://wrong.test", "rel": "self"}]},
+        )
+
+        with pytest.raises(OpenEoClientException, match="No canonical link found in job results metadata"):
+            con120.load_stac_from_job("j0bi6", use_canonical_link="require")
+
+    def test_load_stac_from_job_avoid_canonical(self, con120, requests_mock):
+        results = requests_mock.get(
+            API_URL + "jobs/j0bi6/results",
+            json={"links": [{"href": "https://stac.test", "rel": "canonical"}]},
+        )
+
+        cube = con120.load_stac_from_job("j0bi6", use_canonical_link="avoid")
+
+        assert results.called is False
+        assert cube.flat_graph()["loadstac1"]["arguments"]["url"] == API_URL + "jobs/j0bi6/results"
+
+    def test_load_stac_from_job_invalid_canonical_link_mode(self, con120):
+        with pytest.raises(ValueError, match="Invalid use_canonical_link mode 'sometimes'"):
+            con120.load_stac_from_job("j0bi6", use_canonical_link="sometimes")
+
     def test_load_stac_from_job_from_jobid(self, con120, requests_mock):
         requests_mock.get(
             API_URL + "jobs/j0bi6/results",