Use same requests session in oauth2 token provider as well. (#99)

* Use same requests session in oauth2 token provider as well.

* Improve CC Flow Error Handling

* add more tests and refactor verify behavior

* add docs entry

* feat: release v2.1.17

Author: nicornk

docs/changelog.md

@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog],
 and this project adheres to [Semantic Versioning].
 
+## [2.1.17] - 2025-05-19
+
+## Added 
+- Use same requests session in oauth2 token provider as well. (#99)
+
 ## [2.1.16] - 2025-05-16
 
 ## Added 
@@ -342,6 +347,7 @@ and this project adheres to [Semantic Versioning].
 [Keep a Changelog]: https://keepachangelog.com/en/1.0.0/
 [Semantic Versioning]: https://semver.org/spec/v2.0.0.html
 
+[2.1.17]: https://github.com/emdgroup/foundry-dev-tools/compare/v2.1.16...v2.1.17
 [2.1.16]: https://github.com/emdgroup/foundry-dev-tools/compare/v2.1.15...v2.1.16
 [2.1.15]: https://github.com/emdgroup/foundry-dev-tools/compare/v2.1.14...v2.1.15
 [2.1.14]: https://github.com/emdgroup/foundry-dev-tools/compare/v2.1.13...v2.1.14

docs/configuration.md

@@ -184,6 +184,34 @@ def compute(ctx, output_transform, source: ResolvedSource):
     )
 ```
 
+### Configuration for Python Functions (Serverless)
+
+```
+from functions.api import function
+from functions.sources import get_source
+
+import json
+from foundry_dev_tools import FoundryContext, OAuthTokenProvider, Config
+
+
+@function(sources=["SourceFoundryApi"])
+def get_user_info() -> str:
+    source = get_source("LsRestMTableauExtractSchedulerFoundryApi")
+    client = source.get_https_connection().get_client()
+
+    ctx = FoundryContext(
+        config=Config(requests_session=client),
+        token_provider=OAuthTokenProvider(
+            host="your-stack.palantirfoundry.com",
+            client_id=source.get_secret("ClientID"),
+            client_secret=source.get_secret("ClientSecret"),
+            grant_type="client_credentials",
+        ),
+    )
+    return json.dumps(ctx.multipass.get_user_info())
+
+```
+
 ## How the Configuration Gets Loaded and Merged
 
 For example if there are the files /etc/foundry-dev-tools/config.toml: 

libs/foundry-dev-tools/src/foundry_dev_tools/clients/context_client.py

@@ -4,8 +4,10 @@
 
 import logging
 import numbers
+import os
 import time
 import typing
+from pathlib import Path
 from typing import TYPE_CHECKING
 
 import requests
@@ -72,8 +74,9 @@ class ContextHTTPClient(requests.Session):
 
     def __init__(self, debug: bool = False, requests_ca_bundle: PathLike[str] | None = None) -> None:
         self.debug = debug
-        self.requests_ca_bundle = requests_ca_bundle
         super().__init__()
+        if requests_ca_bundle is not None and Path(requests_ca_bundle).is_file():
+            self.verify = os.fspath(requests_ca_bundle)
 
         self._counter = 0
 
@@ -117,8 +120,6 @@ def request(
             cert: see :py:meth:`requests.Session.request`
             json: see :py:meth:`requests.Session.request`
         """
-        if verify is None and (rcab := self.requests_ca_bundle) is not None:
-            verify = rcab
         if self.debug:
             self._counter = count = self._counter + 1
             LOGGER.debug(f"(r{count}) Making {method!s} request to {url!s}")  # noqa: G004

libs/foundry-dev-tools/src/foundry_dev_tools/config/config.py

@@ -46,14 +46,14 @@ class Config:
 
     def __init__(
         self,
-        requests_ca_bundle: PathLike[str] | None = None,
+        requests_ca_bundle: PathLike[str] | str | None = None,
         transforms_sql_sample_row_limit: int = 5000,
         transforms_sql_dataset_size_threshold: int = 500,
         transforms_sql_sample_select_random: bool = False,
         transforms_force_full_dataset_download: bool = False,
-        cache_dir: PathLike[str] | None = None,
+        cache_dir: PathLike[str] | str | None = None,
         transforms_freeze_cache: bool = False,
-        transforms_output_folder: PathLike[str] | None = None,
+        transforms_output_folder: PathLike[str] | str | None = None,
         rich_traceback: bool = False,
         debug: bool = False,
         requests_session: requests.Session | None = None,

libs/foundry-dev-tools/src/foundry_dev_tools/config/context.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import os
 from functools import cached_property, partial
 from typing import TYPE_CHECKING
 
@@ -69,12 +68,11 @@ def __init__(
         else:
             self.client = self.config.requests_session
 
+        self.token_provider.set_requests_session(self.client)
         self.client.auth = lambda r: self.token_provider.requests_auth_handler(r)
         self.client.headers["User-Agent"] = requests.utils.default_user_agent(
             f"foundry-dev-tools/{__version__}/python-requests"
         )
-        if self.config.requests_ca_bundle:
-            self.verify = os.fspath(self.config.requests_ca_bundle)
 
         if self.config.rich_traceback:
             from rich.traceback import install

libs/foundry-dev-tools/src/foundry_dev_tools/config/token_provider.py

@@ -13,7 +13,8 @@
 
 from foundry_dev_tools.config.config_types import Host
 from foundry_dev_tools.errors.config import TokenProviderConfigError
-from foundry_dev_tools.errors.meta import FoundryAPIError
+from foundry_dev_tools.errors.handling import ErrorHandlingConfig, raise_foundry_api_error
+from foundry_dev_tools.errors.multipass import ClientAuthenticationFailedError
 from foundry_dev_tools.utils.config import entry_point_fdt_token_provider
 
 if TYPE_CHECKING:
@@ -52,6 +53,9 @@ def requests_auth_handler(self, r: requests.PreparedRequest) -> requests.Prepare
         r.headers.setdefault("authorization", f"Bearer {self.token}")
         return r
 
+    def set_requests_session(self, session: requests.Session) -> None:
+        """No-op by default."""
+
 
 class JWTTokenProvider(TokenProvider):
     """Provides Host and Token."""
@@ -150,6 +154,7 @@ def __init__(
                 self.scopes = DEFAULT_OAUTH_SCOPES
         else:
             self.scopes = scopes
+        self._requests_session = requests.Session()
 
     def _scopes_to_list(self, scopes: list[str] | str | None) -> list[str] | None:
         if scopes is not None and isinstance(scopes, str):
@@ -168,27 +173,38 @@ def _request_token(self) -> tuple[Token, float]:
             )
             return credentials.token, credentials.expiry.timestamp()
         if self.grant_type == "client_credentials" and self._client_secret is not None:
-            resp = requests.request(
-                "POST",
-                f"{self.host.url}/multipass/api/oauth2/token",
-                data={"grant_type": "client_credentials", "scope": " ".join(self.scopes)}
-                if self.scopes
-                else {"grant_type": "client_credentials"},
-                headers={
-                    "Content-Type": "application/x-www-form-urlencoded",
-                    "Authorization": "Basic "
-                    + base64.b64encode(
-                        bytes(
-                            self._client_id + ":" + self._client_secret,
-                            "ISO-8859-1",
-                        ),
-                    ).decode("ascii"),
-                },
-                timeout=30,
+            # since we share the same requests session everywhere and on this session
+            # the auth is set to a lambda function we need to manually disable this
+            # for the single token call
+            auth_handler = self._requests_session.auth
+            self._requests_session.auth = None
+            try:
+                resp = self._requests_session.request(
+                    "POST",
+                    f"{self.host.url}/multipass/api/oauth2/token",
+                    data={"grant_type": "client_credentials", "scope": " ".join(self.scopes)}
+                    if self.scopes
+                    else {"grant_type": "client_credentials"},
+                    headers={
+                        "Content-Type": "application/x-www-form-urlencoded",
+                        "Authorization": "Basic "
+                        + base64.b64encode(
+                            bytes(
+                                self._client_id + ":" + self._client_secret,
+                                "ISO-8859-1",
+                            ),
+                        ).decode("ascii"),
+                    },
+                    timeout=30,
+                )
+            finally:
+                # add original auth handler again
+                self._requests_session.auth = auth_handler
+            raise_foundry_api_error(
+                resp,
+                error_handling=ErrorHandlingConfig({401: ClientAuthenticationFailedError}, client_id=self._client_id),
             )
             credentials = resp.json()
-            if "error" in credentials:
-                raise FoundryAPIError(resp)
             return credentials["access_token"], credentials["expires_in"] + time.time()
         if self._client_secret is None:
             msg = f"For grant type {self.grant_type} you need to set a client_secret."
@@ -197,6 +213,10 @@ def _request_token(self) -> tuple[Token, float]:
         msg = f"Grant type {self.grant_type} is not implemented."
         raise NotImplementedError(msg)
 
+    def set_requests_session(self, session: requests.Session) -> None:
+        """Sets request session used for client credentials grant.."""
+        self._requests_session = session
+
 
 class AppServiceTokenProvider(CachedTokenProvider):
     """Token Provider for the AppService, which gets the token via a header from flask/dash/streamlit."""

libs/foundry-dev-tools/src/foundry_dev_tools/errors/multipass.py

@@ -1,9 +1,23 @@
 """Multipass specific errors."""
 
+import requests
+
 from foundry_dev_tools.errors.meta import FoundryAPIError
 
 
 class DuplicateGroupNameError(FoundryAPIError):
     """Exception is thrown when the group name already exists."""
 
     message = "The group name already exists!"
+
+
+class ClientAuthenticationFailedError(FoundryAPIError):
+    """Exception is thrown when client credentials grant failed."""
+
+    def __init__(
+        self, response: requests.Response | None = None, info: str | None = None, client_id: str | None = None
+    ):
+        self.client_id = client_id
+        self.message = "Client authentication failed (invalid_client)."
+
+        super().__init__(response=response, info=info, client_id=client_id)

tests/unit/clients/test_context_client.py

@@ -1,12 +1,16 @@
+import os
 import time
+from pathlib import Path
 from unittest import mock
 
+import pytest
 import requests
 from requests_mock import ANY
 
+from foundry_dev_tools import Config
 from foundry_dev_tools.__about__ import __version__
-from foundry_dev_tools.clients.context_client import DEFAULT_TIMEOUT
-from tests.unit.mocks import MockOAuthTokenProvider
+from foundry_dev_tools.clients.context_client import DEFAULT_TIMEOUT, ContextHTTPClient
+from tests.unit.mocks import FoundryMockContext, MockOAuthTokenProvider, MockTokenProvider
 
 
 def test_context_http_client(test_context_mock, foundry_client_id):
@@ -51,3 +55,41 @@ def test_retry_on_connection_error(test_context_mock):
         m.side_effect = [requests.exceptions.ConnectionError(), response_mock]
         response = test_context_mock.client.request("GET", "test_call_args")
         assert response is response_mock
+
+
+def test_requests_ca_bundle(request, tmp_path_factory):
+    cert_dir = tmp_path_factory.mktemp(f"foundry_dev_tools_test__{request.node.name}").absolute()
+    with Path.open(cert_dir / "ca-bundle.pem", "w") as f:
+        f.write("test")
+    client1 = FoundryMockContext(
+        Config(
+            requests_ca_bundle=os.fspath(cert_dir / "ca-bundle.pem"),
+        ),
+        MockTokenProvider(jwt=request.node.name + "_token"),
+    ).client
+    assert client1.verify == os.fspath(cert_dir / "ca-bundle.pem")
+    client_with_pathlib = FoundryMockContext(
+        Config(
+            requests_ca_bundle=cert_dir / "ca-bundle.pem",
+        ),
+        MockTokenProvider(jwt=request.node.name + "_token"),
+    ).client
+    assert client_with_pathlib.verify == os.fspath(cert_dir / "ca-bundle.pem")
+
+    client = ContextHTTPClient(debug=False, requests_ca_bundle=None)
+    assert client.verify is True
+
+    with pytest.raises(TypeError):
+        _ = ContextHTTPClient(debug=False, requests_ca_bundle=False)
+
+    session = requests.Session()
+    assert session.verify is True
+
+
+def test_verify_is_passed(test_context_mock):
+    client1 = test_context_mock.client
+    assert client1.verify is True
+    test_context_mock.mock_adapter.register_uri(ANY, ANY)
+
+    req = test_context_mock.client.request("POST", "http+mock://test_authorization", verify=False).request
+    assert req.verify is False

tests/unit/config/test_context.py

@@ -6,10 +6,12 @@
 from unittest import mock
 
 import pytest
+import requests_mock
 
-from foundry_dev_tools import FoundryContext
+from foundry_dev_tools import Config, FoundryContext
 from foundry_dev_tools.errors.config import FoundryConfigError
-from tests.unit.mocks import TEST_DOMAIN
+from foundry_dev_tools.errors.multipass import ClientAuthenticationFailedError
+from tests.unit.mocks import TEST_DOMAIN, FoundryMockContext, MockOAuthTokenProvider
 
 if TYPE_CHECKING:
     from pathlib import Path
@@ -26,3 +28,27 @@ def test_app_service(mock_config_location: dict[Path, None]):
     ):
         warnings.simplefilter("error")  # raise errors on warnings
         _ = FoundryContext()
+
+
+def test_client_credentials_error_handling(foundry_token):
+    ctx = FoundryMockContext(
+        config=Config(),
+        token_provider=MockOAuthTokenProvider(
+            client_id="test_client_id",
+            client_secret="test",  # noqa: S106
+            grant_type="client_credentials",
+        ),
+    )
+    with requests_mock.Mocker() as m:
+        m.post(
+            f"{ctx.token_provider.host.url}/multipass/api/oauth2/token",
+            json={"error": "invalid_client", "error_description": "Client authentication failed"},
+            status_code=401,
+        )
+        m.get(
+            f"{ctx.token_provider.host.url}/multipass/api/me",
+            json={"id": "id", "attributes": "a", "username": "u"},
+        )
+        with pytest.raises(ClientAuthenticationFailedError) as catched:
+            ctx.get_user_info()
+    assert catched.value.client_id == "test_client_id"