Add option to keep Authorization header on requests

TM1py/Services/RestService.py

@@ -107,6 +107,7 @@ def __init__(self, **kwargs):
         - **namespace** (str): Optional CAM namespace.
         - **cam_passport** (str): The CAM passport.
         - **session_id** (str): TM1SessionId, e.g. q7O6e1w49AixeuLVxJ1GZg.
+        - **access_token** (str|callable): Bearer token, or a zero-argument callable returning one. A callable is invoked on every connect and reconnect, so a caller that refreshes tokens (e.g. an OAuth client) supplies a fresh token after a session timeout.
         - **application_client_id** (str): Planning Analytics Engine (v12) named application client ID created via manage service.
         - **application_client_secret** (str): Planning Analytics Engine (v12) named application secret created via manage service.
         - **api_key** (str): Planning Analytics Engine (v12) API Key from https://cloud.ibm.com/iam/apikeys.
@@ -143,6 +144,7 @@ def __init__(self, **kwargs):
         - **compress_request_body** (bool): Gzip-compress request bodies at the transport seam. Opt-in. Default: False.
         - **gzip_min_bytes** (int): Minimum (encoded) request body size in bytes to compress. Default: 1024.
         - **gzip_compress_level** (int): Gzip compression level, 1 (fastest) to 9 (smallest). Default: 6.
+        - **keep_authorization_header** (bool): Keep sending the Authorization header after the session starts, instead of relying on the session cookie alone. Only for token-based gateways that validate the token on every request, such as Planning Analytics Workspace's OAuth route. Default: False.
 
         :param kwargs: See description above for all supported arguments
         """
@@ -172,6 +174,7 @@ def __init__(self, **kwargs):
         self._connection_pool_size = int(kwargs.get("connection_pool_size", self.DEFAULT_CONNECTION_POOL_SIZE))
         self._pool_connections = int(kwargs.get("pool_connections", self.DEFAULT_POOL_CONNECTIONS))
         self._re_connect_on_session_timeout = kwargs.get("re_connect_on_session_timeout", True)
+        self._keep_authorization_header = self.translate_to_boolean(kwargs.get("keep_authorization_header", False))
         self._re_connect_on_remote_disconnect = kwargs.get("re_connect_on_remote_disconnect", True)
         self._remote_disconnect_max_retries = int(kwargs.get("remote_disconnect_max_retries", 5))
         self._remote_disconnect_retry_delay = float(kwargs.get("remote_disconnect_retry_delay", 1))
@@ -918,6 +921,21 @@ def _extract_tm1_session_id_from_set_cookie_header(auth_response_headers: object
         else:
             return None
 
+    def _resolve_access_token(self):
+        """Return the current access token.
+
+        Supports ``access_token`` being either a plain string or a zero-argument
+        callable. A callable is invoked on every connect and reconnect, so a caller
+        that refreshes tokens (e.g. an OAuth session) can supply a fresh token when
+        TM1py re-establishes a timed-out session, rather than reusing a stale one.
+        """
+        access_token = self._kwargs.get("access_token")
+        if callable(access_token):
+            access_token = access_token()
+        if not access_token:
+            raise ValueError("'access_token' did not resolve to a non-empty token")
+        return access_token
+
     def _start_session(
         self,
         user: str,
@@ -959,7 +977,7 @@ def _start_session(
             self.add_http_header("Authorization", "Bearer " + access_token)
 
         elif self._auth_mode == AuthenticationMode.ACCESS_TOKEN:
-            self.add_http_header("Authorization", "Bearer " + self._kwargs.get("access_token"))
+            self.add_http_header("Authorization", "Bearer " + self._resolve_access_token())
 
         # v11 authorization (Basic, CAM) through Headers
         else:
@@ -1040,8 +1058,12 @@ def _start_session(
             if session_id is not None:
                 self._s.cookies.set("TM1SessionId", session_id)
 
-            # After we have session cookie, drop the Authorization Header
-            self.remove_http_header("Authorization")
+            # After we have session cookie, drop the Authorization Header, unless the
+            # server authenticates every request rather than only the session-start
+            # request (e.g. an OAuth gateway such as Planning Analytics Workspace's
+            # /api/.../tm1/<database> route). keep_authorization_header leaves it on.
+            if not self._keep_authorization_header:
+                self.remove_http_header("Authorization")
 
     def _url_and_body(self, url: str, data: str, encoding: str = "utf-8") -> Tuple[str, bytes]:
         """create proper url and payload"""

Tests/RestService_authorization_header_test.py

@@ -0,0 +1,79 @@
+import unittest
+from unittest.mock import patch
+
+import requests
+
+from TM1py.Services.RestService import AuthenticationMode, RestService
+
+BASE_URL = "https://tm1.example.com/api/v1/tm1/CONS"
+
+
+def _ok_response():
+    response = requests.Response()
+    response.status_code = 200
+    response._content = b"12.0.0"
+    return response
+
+
+class TestRestServiceAuthorizationHeader(unittest.TestCase):
+    """Bearer-token session start: keep_authorization_header and callable access_token.
+
+    Builds a real RestService with connect() and set_version() patched out, then runs the real
+    _start_session against a stubbed HTTP GET, so no TM1 server is needed.
+    """
+
+    def _rest(self, **kwargs):
+        with patch.object(RestService, "connect"), patch.object(RestService, "set_version"):
+            rest = RestService(base_url=BASE_URL, **kwargs)
+        self.assertEqual(rest._auth_mode, AuthenticationMode.ACCESS_TOKEN)
+        self.sent_authorization = []
+
+        def get(url, headers, **_):
+            self.sent_authorization.append(headers.get("Authorization"))
+            return _ok_response()
+
+        rest._s.get = get
+        return rest
+
+    def _start_session(self, rest):
+        rest._start_session(user=None, password=None)
+
+    def test_header_removed_by_default(self):
+        rest = self._rest(access_token="abc")
+        self._start_session(rest)
+
+        self.assertEqual(self.sent_authorization, ["Bearer abc"])
+        self.assertNotIn("Authorization", rest._headers)
+
+    def test_header_kept_when_flag_on(self):
+        rest = self._rest(access_token="abc", keep_authorization_header=True)
+        self._start_session(rest)
+
+        self.assertEqual(rest._headers["Authorization"], "Bearer abc")
+
+    def test_flag_given_as_false_string_stays_off(self):
+        """config.ini values arrive as strings; "False" must not switch the option on."""
+        rest = self._rest(access_token="abc", keep_authorization_header="False")
+        self._start_session(rest)
+
+        self.assertNotIn("Authorization", rest._headers)
+
+    def test_callable_token_is_resolved_on_every_session_start(self):
+        tokens = iter(["first", "second"])
+        rest = self._rest(access_token=lambda: next(tokens), keep_authorization_header=True)
+
+        self._start_session(rest)
+        self._start_session(rest)
+
+        self.assertEqual(self.sent_authorization, ["Bearer first", "Bearer second"])
+        self.assertEqual(rest._headers["Authorization"], "Bearer second")
+
+    def test_callable_returning_empty_token_raises(self):
+        rest = self._rest(access_token=lambda: "")
+
+        with self.assertRaises(ValueError):
+            self._start_session(rest)
+
+
+if __name__ == "__main__":
+    unittest.main()