Issue #560 Add on_response_headers to DataCube.download() et al.

Author: soxofaan

CHANGELOG.md

@@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `sar_backscatter`: try to retrieve coefficient options from backend ([#693](https://github.com/Open-EO/openeo-python-client/issues/693))
 - Improve error message when OIDC provider is unavailable ([#751](https://github.com/Open-EO/openeo-python-client/issues/751))
+- Added `on_response_headers` argument to `DataCube.download()` and related to handle (e.g. `print`) the response headers ([#560](https://github.com/Open-EO/openeo-python-client/issues/560))
 
 ### Changed
 

openeo/rest/_testing.py

@@ -179,6 +179,7 @@ def _handle_post_result(self, request, context):
         elif isinstance(result, str):
             result = result.encode("utf-8")
         assert isinstance(result, bytes)
+        context.headers["OpenEO-Identifier"] = f"r-{len(self.sync_requests):03d}"
         return result
 
     def _handle_post_jobs(self, request, context):

openeo/rest/connection.py

@@ -19,6 +19,7 @@
     Iterable,
     Iterator,
     List,
+    Mapping,
     Optional,
     Sequence,
     Set,
@@ -1568,6 +1569,7 @@ def download(
         chunk_size: int = DEFAULT_DOWNLOAD_CHUNK_SIZE,
         additional: Optional[dict] = None,
         job_options: Optional[dict] = None,
+        on_response_headers: Optional[Callable[[Mapping], None]] = None,
     ) -> Union[None, bytes]:
         """
         Send the underlying process graph to the backend
@@ -1586,13 +1588,17 @@ def download(
         :param additional: (optional) additional (top-level) properties to set in the request body
         :param job_options: (optional) dictionary of job options to pass to the backend
             (under top-level property "job_options")
+        :param on_response_headers: (optional) callback to handle/show the response headers
 
         :return: if ``outputfile`` was not specified:
             a :py:class:`bytes` object containing the raw data.
             Otherwise, ``None`` is returned.
 
-        .. versionadded:: 0.36.0
+        .. versionchanged:: 0.36.0
             Added arguments ``additional`` and ``job_options``.
+
+        .. versionchanged:: 0.40
+            Added argument ``on_response_headers``.
         """
         pg_with_metadata = self._build_request_with_process_graph(
             process_graph=graph, additional=additional, job_options=job_options
@@ -1605,6 +1611,8 @@ def download(
             stream=True,
             timeout=timeout or DEFAULT_TIMEOUT_SYNCHRONOUS_EXECUTE,
         )
+        if on_response_headers:
+            on_response_headers(response.headers)
 
         if outputfile is not None:
             target = Path(outputfile)

openeo/rest/datacube.py

@@ -17,7 +17,18 @@
 import urllib.parse
 import warnings
 from builtins import staticmethod
-from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
 
 import numpy as np
 import requests
@@ -2386,6 +2397,7 @@ def download(
         auto_add_save_result: bool = True,
         additional: Optional[dict] = None,
         job_options: Optional[dict] = None,
+        on_response_headers: Optional[Callable[[Mapping], None]] = None,
     ) -> Union[None, bytes]:
         """
         Send the underlying process graph to the backend
@@ -2403,6 +2415,7 @@ def download(
         :param additional: (optional) additional (top-level) properties to set in the request body
         :param job_options: (optional) dictionary of job options to pass to the backend
             (under top-level property "job_options")
+        :param on_response_headers: (optional) callback to handle/show the response headers
 
         :return: if ``outputfile`` was not specified:
             a :py:class:`bytes` object containing the raw data.
@@ -2413,14 +2426,22 @@ def download(
 
         .. versionchanged:: 0.36.0
             Added arguments ``additional`` and ``job_options``.
+
+        .. versionchanged:: 0.40
+            Added argument ``on_response_headers``.
         """
         # TODO #278 centralize download/create_job/execute_job logic in DataCube, VectorCube, MlModel, ...
         if auto_add_save_result:
             res = self._auto_save_result(format=format, outputfile=outputfile, options=options)
         else:
             res = self
         return self._connection.download(
-            res.flat_graph(), outputfile=outputfile, validate=validate, additional=additional, job_options=job_options
+            res.flat_graph(),
+            outputfile=outputfile,
+            validate=validate,
+            additional=additional,
+            job_options=job_options,
+            on_response_headers=on_response_headers,
         )
 
     def validate(self) -> List[dict]:

openeo/rest/stac_resource.py

@@ -2,7 +2,7 @@
 
 import typing
 from pathlib import Path
-from typing import Optional, Union
+from typing import Callable, Mapping, Optional, Union
 
 from openeo.internal.documentation import openeo_process
 from openeo.internal.graph_building import PGNode
@@ -87,6 +87,7 @@ def download(
         validate: Optional[bool] = None,
         additional: Optional[dict] = None,
         job_options: Optional[dict] = None,
+        on_response_headers: Optional[Callable[[Mapping], None]] = None,
     ):
         """
         Send the underlying process graph to the backend
@@ -101,17 +102,22 @@ def download(
         :param additional: (optional) additional (top-level) properties to set in the request body
         :param job_options: (optional) dictionary of job options to pass to the backend
             (under top-level property "job_options")
+        :param on_response_headers: (optional) callback to handle/show the response headers
 
         :return: if ``outputfile`` was not specified:
             a :py:class:`bytes` object containing the raw data.
             Otherwise, ``None`` is returned.
+
+        .. versionchanged:: 0.40
+            Added argument ``on_response_headers``.
         """
         return self._connection.download(
             graph=self.flat_graph(),
             outputfile=outputfile,
             validate=validate,
             additional=additional,
             job_options=job_options,
+            on_response_headers=on_response_headers,
         )
 
     def create_job(

openeo/rest/vectorcube.py

@@ -4,7 +4,7 @@
 import logging
 import pathlib
 import typing
-from typing import Callable, List, Optional, Tuple, Union
+from typing import Callable, List, Mapping, Optional, Tuple, Union
 
 import shapely.geometry.base
 
@@ -250,6 +250,7 @@ def download(
         auto_add_save_result: bool = True,
         additional: Optional[dict] = None,
         job_options: Optional[dict] = None,
+        on_response_headers: Optional[Callable[[Mapping], None]] = None,
     ) -> Union[None, bytes]:
         """
         Send the underlying process graph to the backend
@@ -267,6 +268,7 @@ def download(
         :param additional: (optional) additional (top-level) properties to set in the request body
         :param job_options: (optional) dictionary of job options to pass to the backend
             (under top-level property "job_options")
+        :param on_response_headers: (optional) callback to handle/show the response headers
 
         :return: if ``outputfile`` was not specified:
             a :py:class:`bytes` object containing the raw data.
@@ -280,6 +282,9 @@ def download(
 
         .. versionchanged:: 0.39.0
             Added arguments ``additional`` and ``job_options``.
+
+        .. versionchanged:: 0.40
+            Added argument ``on_response_headers``.
         """
         # TODO #278 centralize download/create_job/execute_job logic in DataCube, VectorCube, MlModel, ...
         if auto_add_save_result:

tests/rest/test_connection.py

@@ -3612,6 +3612,16 @@ def test_download_with_additional_and_job_options(dummy_backend, tmp_path):
     }
 
 
+def test_download_on_response_headers(dummy_backend, tmp_path):
+    results = []
+    dummy_backend.connection.download(
+        {"foo1": {"process_id": "foo"}},
+        tmp_path / "result.data",
+        on_response_headers=results.append,
+    )
+    assert results == [{"OpenEO-Identifier": "r-001"}]
+
+
 @pytest.mark.parametrize(
     "pg",
     [