Add base headers in properties to signer_headers.

Author: iamthomaspowell

pyiceberg/catalog/rest.py

@@ -79,7 +79,7 @@
 from pyiceberg.typedef import EMPTY_DICT, UTF8, IcebergBaseModel, Identifier, Properties
 from pyiceberg.types import transform_dict_value_to_str
 from pyiceberg.utils.deprecated import deprecation_message
-from pyiceberg.utils.properties import get_first_property_value, property_as_bool
+from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool
 
 if TYPE_CHECKING:
     import pyarrow as pa
@@ -139,7 +139,6 @@ class IdentifierKind(Enum):
 SIGV4_SERVICE = "rest.signing-name"
 AUTH_URL = "rest.authorization-url"
 OAUTH2_SERVER_URI = "oauth2-server-uri"
-HEADER_PREFIX = "header."
 
 NAMESPACE_SEPARATOR = b"\x1f".decode(UTF8)
 
@@ -554,16 +553,13 @@ def _refresh_token(self, session: Optional[Session] = None, initial_token: Optio
             session.headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {token}"
 
     def _config_headers(self, session: Session) -> None:
-        header_properties = self._extract_headers_from_properties()
+        header_properties = get_header_properties(self.properties)
         session.headers.update(header_properties)
         session.headers["Content-type"] = "application/json"
         session.headers["X-Client-Version"] = ICEBERG_REST_SPEC_VERSION
         session.headers["User-Agent"] = f"PyIceberg/{__version__}"
         session.headers.setdefault("X-Iceberg-Access-Delegation", ACCESS_DELEGATION_DEFAULT)
 
-    def _extract_headers_from_properties(self) -> Dict[str, str]:
-        return {key[len(HEADER_PREFIX) :]: value for key, value in self.properties.items() if key.startswith(HEADER_PREFIX)}
-
     def _create_table(
         self,
         identifier: Union[str, Identifier],

pyiceberg/io/fsspec.py

@@ -79,7 +79,7 @@
     OutputStream,
 )
 from pyiceberg.typedef import Properties
-from pyiceberg.utils.properties import get_first_property_value, property_as_bool
+from pyiceberg.utils.properties import get_first_property_value, get_header_properties, property_as_bool
 
 logger = logging.getLogger(__name__)
 
@@ -91,6 +91,7 @@ def s3v4_rest_signer(properties: Properties, request: AWSRequest, **_: Any) -> A
     signer_headers = {}
     if token := properties.get(TOKEN):
         signer_headers = {"Authorization": f"Bearer {token}"}
+    signer_headers.update(get_header_properties(properties))
 
     signer_body = {
         "method": request.method,

pyiceberg/utils/properties.py

@@ -25,6 +25,9 @@
 from pyiceberg.types import strtobool
 
 
+HEADER_PREFIX = "header."
+
+
 def property_as_int(
     properties: Dict[str, str],
     property_name: str,
@@ -74,3 +77,13 @@ def get_first_property_value(
         if property_value := properties.get(property_name):
             return property_value
     return None
+
+
+def get_header_properties(
+    properties: Properties,
+) -> Properties:
+    return {
+        key[len(HEADER_PREFIX):]: value
+        for key, value in properties.items()
+        if key.startswith(HEADER_PREFIX)
+    }

tests/io/test_fsspec.py

@@ -692,6 +692,7 @@ def test_s3v4_rest_signer(requests_mock: Mocker) -> None:
                 "X-Amz-Security-Token": [
                     "YQoJb3JpZ2luX2VjEDoaCXVzLXdlc3QtMiJGMEQCID/fFxZP5oaEgQmcwP6XhZa0xSq9lmLSx8ffaWbySfUPAiAesa7sjd/WV4uwRTO0S03y/MWVtgpH+/NyZQ4bZgLVriqrAggTEAEaDDAzMzQwNzIyMjE1OSIMOeFOWhZIurMmAqjsKogCxMCqxX8ZjK0gacAkcDqBCyA7qTSLhdfKQIH/w7WpLBU1km+cRUWWCudan6gZsAq867DBaKEP7qI05DAWr9MChAkgUgyI8/G3Z23ET0gAedf3GsJbakB0F1kklx8jPmj4BPCht9RcTiXiJ5DxTS/cRCcalIQXmPFbaJSqpBusVG2EkWnm1v7VQrNPE2Os2b2P293vpbhwkyCEQiGRVva4Sw9D1sKvqSsK10QCRG+os6dFEOu1kARaXi6pStvR4OVmj7OYeAYjzaFchn7nz2CSae0M4IluiYQ01eQAywbfRo9DpKSmDM/DnPZWJnD/woLhaaaCrCxSSEaFsvGOHFhLd3Rknw1v0jADMILUtJoGOp4BpqKqyMz0CY3kpKL0jfR3ykTf/ge9wWVE0Alr7wRIkGCIURkhslGHqSyFRGoTqIXaxU+oPbwlw/0w/nYO7qQ6bTANOWye/wgw4h/NmJ6vU7wnZTXwREf1r6MF72++bE/fMk19LfVb8jN/qrUqAUXTc8gBAUxL5pgy8+oT/JnI2BkVrrLS4ilxEXP9Ahm+6GDUYXV4fBpqpZwdkzQ/5Gw="
                 ],
+                "X-Custom-Header": ["value"],
             },
             "extensions": {},
         },
@@ -714,7 +715,7 @@ def test_s3v4_rest_signer(requests_mock: Mocker) -> None:
         "retries": {"attempt": 1, "invocation-id": "75d143fb-0219-439b-872c-18213d1c8d54"},
     }
 
-    signed_request = s3v4_rest_signer({"token": "abc", "uri": TEST_URI}, request)
+    signed_request = s3v4_rest_signer({"token": "abc", "uri": TEST_URI, "header.X-Custom-Header": "value"}, request)
 
     assert signed_request.url == new_uri
     assert dict(signed_request.headers) == {