Control canonical links in load_stac_from_job

CHANGELOG.md

@@ -9,6 +9,7 @@ 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))
 - `DataCube.resample_spatial()` now supports parameterized `resolution` and `projection` arguments.  ([#897](https://github.com/Open-EO/openeo-python-client/issues/897))
 - Sanitize asset download filenames (e.g. strip slashes, (semi)colon, hash, ...), instead of blindly using the asset key as filename. ([#820](https://github.com/Open-EO/openeo-python-client/issues/820))
 - Support parameters in `DataCube` apply- and band-math operations ([#903](https://github.com/Open-EO/openeo-python-client/issues/903))

openeo/rest/connection.py

@@ -19,6 +19,7 @@
     Iterable,
     Iterator,
     List,
+    Literal,
     Mapping,
     Optional,
     Sequence,
@@ -1470,28 +1471,49 @@ 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,
+        *,
+        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 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.51.0
+            Added the ``canonical_link`` argument.
         """
-        # TODO #634 add option to require or avoid the canonical link
+        if canonical_link not in {"auto", "require", "avoid"}:
+            raise ValueError(f"Invalid canonical_link mode {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")
 
+        stac_link = job.get_results_metadata_url(full=True)
+        if canonical_link == "avoid":
+            return self.load_stac(
+                url=stac_link,
+                spatial_extent=spatial_extent,
+                temporal_extent=temporal_extent,
+                bands=bands,
+                properties=properties,
+            )
+
         try:
             job_results = job.get_results()
 
@@ -1501,17 +1523,20 @@ def load_stac_from_job(
                 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)
+                message = "No canonical link found in job results metadata."
+                if 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 canonical_link == "require":
+                raise
             _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)
 
         return self.load_stac(
             url=stac_link,

tests/rest/test_connection.py

@@ -3240,7 +3240,8 @@ def test_property_filtering(self, con120, properties, expected):
             }
         }
 
-    def test_load_stac_from_job_canonical(self, con120, requests_mock):
+    @pytest.mark.parametrize("canonical_link", ["auto", "require"])
+    def test_load_stac_from_job_canonical(self, con120, requests_mock, canonical_link):
         requests_mock.get(
             API_URL + "jobs/j0bi6/results",
             json={
@@ -3257,7 +3258,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, canonical_link=canonical_link)
 
         fg = cube.flat_graph()
         assert fg == {
@@ -3294,6 +3295,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", 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", 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 canonical_link mode 'sometimes'"):
+            con120.load_stac_from_job("j0bi6", canonical_link="sometimes")
+
     def test_load_stac_from_job_from_jobid(self, con120, requests_mock):
         requests_mock.get(
             API_URL + "jobs/j0bi6/results",