Python: Add the REST token to the properties

Author: Fokko

python/pyiceberg/catalog/__init__.py

@@ -40,6 +40,7 @@
 
 _ENV_CONFIG = Config()
 
+TOKEN = "token"
 TYPE = "type"
 URI = "uri"
 

python/pyiceberg/catalog/rest.py

@@ -34,7 +34,7 @@
     Catalog,
     Identifier,
     Properties,
-    PropertiesUpdateSummary,
+    PropertiesUpdateSummary, TOKEN,
 )
 from pyiceberg.exceptions import (
     AuthorizationExpiredError,
@@ -87,7 +87,6 @@ class Endpoints:
 CREDENTIAL = "credential"
 GRANT_TYPE = "grant_type"
 SCOPE = "scope"
-TOKEN = "token"
 TOKEN_EXCHANGE = "urn:ietf:params:oauth:grant-type:token-exchange"
 SEMICOLON = ":"
 KEY = "key"
@@ -206,12 +205,13 @@ def _create_session(self) -> None:
                 elif ssl_client_cert := ssl_client.get(CERT):
                     self.session.cert = ssl_client_cert
 
+        # If we have credentials, but not a token, we want to fetch a token
+        if TOKEN not in self.properties and CREDENTIAL in self.properties:
+            self.properties[TOKEN] = self._fetch_access_token(self.properties[CREDENTIAL])
+
         # Set Auth token for subsequent calls in the session
         if token := self.properties.get(TOKEN):
             self.session.headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {token}"
-        elif credential := self.properties.get(CREDENTIAL):
-            token = self._fetch_access_token(credential)
-            self.session.headers[AUTHORIZATION_HEADER] = f"{BEARER_PREFIX} {token}"
 
         # Set HTTP headers
         self.session.headers["Content-type"] = "application/json"

python/pyiceberg/io/fsspec.py

@@ -27,6 +27,7 @@
 from requests import HTTPError
 from s3fs import S3FileSystem
 
+from pyiceberg.catalog import TOKEN
 from pyiceberg.exceptions import SignError
 from pyiceberg.io import FileIO, InputFile, OutputFile
 from pyiceberg.typedef import Properties
@@ -35,16 +36,20 @@
 
 
 def s3v4_rest_signer(properties: Properties, request: AWSRequest, **_) -> AWSRequest:
+    if TOKEN not in properties:
+        raise SignError("Signer set, but token is not available")
+
     signer_url = properties["uri"].rstrip("/")
-    signer_headers = {"Authorization": f"Bearer {properties['token']}"}
+    signer_headers = {"Authorization": f"Bearer {properties[TOKEN]}"}
     signer_body = {
         "method": request.method,
         "region": request.context["client_region"],
         "uri": request.url,
         "headers": {key: [val] for key, val in request.headers.items()},
     }
+
+    response = requests.post(f"{signer_url}/v1/aws/s3/sign", headers=signer_headers, json=signer_body)
     try:
-        response = requests.post(f"{signer_url}/v1/aws/s3/sign", headers=signer_headers, json=signer_body)
         response.raise_for_status()
         response_json = response.json()
     except HTTPError as e: