Inital version for oauth_client support in client

Author: vaibhavatlan

pyatlan/client/aio/client.py

@@ -90,6 +90,7 @@ class AsyncAtlanClient(AtlanClient):
     """
 
     _async_session: Optional[httpx.AsyncClient] = PrivateAttr(default=None)
+    _async_oauth_token_manager: Optional[Any] = PrivateAttr(default=None)
     _async_admin_client: Optional[AsyncAdminClient] = PrivateAttr(default=None)
     _async_asset_client: Optional[AsyncAssetClient] = PrivateAttr(default=None)
     _async_audit_client: Optional[AsyncAuditClient] = PrivateAttr(default=None)
@@ -131,13 +132,19 @@ class AsyncAtlanClient(AtlanClient):
     _async_user_cache: Optional[AsyncUserCache] = PrivateAttr(default=None)
 
     def __init__(self, **kwargs):
-        # Initialize sync client (handles all validation, env vars, etc.)
         super().__init__(**kwargs)
 
-        # Build proxy/SSL configuration (reuse from sync client)
+        if self.oauth_client_id and self.oauth_client_secret:
+            from pyatlan.client.aio.oauth import AsyncOAuthTokenManager
+
+            self._async_oauth_token_manager = AsyncOAuthTokenManager(
+                base_url=str(self.base_url),
+                client_id=self.oauth_client_id,
+                client_secret=self.oauth_client_secret,
+            )
+
         transport_kwargs = self._build_transport_proxy_config(kwargs)
 
-        # Create async session with custom transport that supports retry and proxy
         self._async_session = httpx.AsyncClient(
             transport=PyatlanAsyncTransport(retry=self.retry, **transport_kwargs),
             headers={
@@ -434,17 +441,16 @@ def _api_logger(self, api, path):
     async def _create_params(
         self, api, query_params, request_obj, exclude_unset: bool = True
     ):
-        """
-        Async version of _create_params that uses AsyncAtlanRequest for AtlanObject instances.
-        """
         params = copy.deepcopy(self._request_params)
+        if self._async_oauth_token_manager:
+            token = await self._async_oauth_token_manager.get_token()
+            params["headers"]["authorization"] = f"Bearer {token}"
         params["headers"]["Accept"] = api.consumes
         params["headers"]["content-type"] = api.produces
         if query_params is not None:
             params["params"] = query_params
         if request_obj is not None:
             if isinstance(request_obj, AtlanObject):
-                # Use AsyncAtlanRequest for async retranslation
                 async_request = AsyncAtlanRequest(instance=request_obj, client=self)
                 params["data"] = await async_request.json()
             elif api.consumes == APPLICATION_ENCODED_FORM:
@@ -685,9 +691,8 @@ async def _handle_error_response(
                 "\n".join(error_cause_details) if error_cause_details else ""
             )
 
-            # Retry with impersonation (if _user_id is present) on authentication failure
             if (
-                self._user_id
+                (self._user_id or self._async_oauth_token_manager)
                 and not self._401_has_retried.get()
                 and response.status_code
                 == ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
@@ -742,12 +747,22 @@ async def _handle_401_token_refresh(
         download_file_path=None,
         text_response=False,
     ):
-        """
-        Async version of token refresh and retry logic.
-        Handles token refresh and retries the API request upon a 401 Unauthorized response.
-        """
+        if self._async_oauth_token_manager:
+            await self._async_oauth_token_manager.invalidate_token()
+            token = await self._async_oauth_token_manager.get_token()
+            params["headers"]["authorization"] = f"Bearer {token}"
+            self._401_has_retried.set(True)
+            LOGGER.debug("Successfully refreshed OAuth token after 401.")
+            return await self._call_api_internal(
+                api,
+                path,
+                params,
+                binary_data=binary_data,
+                download_file_path=download_file_path,
+                text_response=text_response,
+            )
+
         try:
-            # Use sync impersonation call since it's a quick API call
             new_token = await self.impersonate.user(user_id=self._user_id)
         except Exception as e:
             LOGGER.debug(
@@ -763,11 +778,9 @@ async def _handle_401_token_refresh(
         self._request_params["headers"]["authorization"] = f"Bearer {self.api_key}"
         LOGGER.debug("Successfully completed async 401 automatic token refresh.")
 
-        # Async retry loop to ensure token is active before retrying original request
         retry_count = 1
         while retry_count <= self.retry.total:
             try:
-                # Use async typedef call to validate token
                 response = await self.typedef.get(
                     type_category=[AtlanTypeCategory.STRUCT]
                 )
@@ -778,10 +791,9 @@ async def _handle_401_token_refresh(
                     "Retrying async to get typedefs (to ensure token is active) after token refresh failed: %s",
                     e,
                 )
-            await asyncio.sleep(retry_count)  # Linear backoff with async sleep
+            await asyncio.sleep(retry_count)
             retry_count += 1
 
-        # Retry the API call with the new token
         return await self._call_api_internal(
             api,
             path,

pyatlan/client/aio/oauth.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+import asyncio
+import time
+from typing import Optional
+
+import httpx
+from authlib.oauth2.rfc6749 import OAuth2Token
+
+
+class AsyncOAuthTokenManager:
+    def __init__(
+        self,
+        base_url: str,
+        client_id: str,
+        client_secret: str,
+        http_client: Optional[httpx.AsyncClient] = None,
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.token_url = f"{self.base_url}/api/service/oauth-clients/token"
+        self._lock = asyncio.Lock()
+        self._http_client = http_client or httpx.AsyncClient(timeout=30.0)
+        self._token: Optional[OAuth2Token] = None
+        self._owns_client = http_client is None
+
+    async def get_token(self) -> str:
+        async with self._lock:
+            if self._token and not self._token.is_expired():
+                return self._token["access_token"]
+
+            response = await self._http_client.post(
+                self.token_url,
+                json={
+                    "clientId": self.client_id,
+                    "clientSecret": self.client_secret,
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
+            data = response.json()
+            access_token = data.get("accessToken") or data.get("access_token")
+            expires_in = data.get("expiresIn") or data.get("expires_in", 600)
+
+            self._token = OAuth2Token(
+                {
+                    "access_token": access_token,
+                    "token_type": data.get("tokenType")
+                    or data.get("token_type", "Bearer"),
+                    "expires_in": expires_in,
+                    "expires_at": int(time.time()) + expires_in,
+                }
+            )
+
+            return access_token
+
+    async def invalidate_token(self):
+        async with self._lock:
+            self._token = None
+
+    async def aclose(self):
+        if self._owns_client:
+            await self._http_client.aclose()

pyatlan/client/atlan.py

@@ -127,7 +127,9 @@ def log_response(response, *args, **kwargs):
 
 class AtlanClient(BaseSettings):
     base_url: Union[Literal["INTERNAL"], HttpUrl]
-    api_key: str
+    api_key: Optional[str] = None
+    oauth_client_id: Optional[str] = None
+    oauth_client_secret: Optional[str] = None
     connect_timeout: float = 30.0  # 30 secs
     read_timeout: float = 900.0  # 15 mins
     retry: Retry = DEFAULT_RETRY
@@ -137,6 +139,7 @@ class AtlanClient(BaseSettings):
     _session: httpx.Client = PrivateAttr()
     _request_params: dict = PrivateAttr()
     _user_id: Optional[str] = PrivateAttr(default=None)
+    _oauth_token_manager: Optional[Any] = PrivateAttr(default=None)
     _workflow_client: Optional[WorkflowClient] = PrivateAttr(default=None)
     _credential_client: Optional[CredentialClient] = PrivateAttr(default=None)
     _admin_client: Optional[AdminClient] = PrivateAttr(default=None)
@@ -172,17 +175,24 @@ class Config:
 
     def __init__(self, **data):
         super().__init__(**data)
-        self._request_params = (
-            {"headers": {"authorization": f"Bearer {self.api_key}"}}
-            if self.api_key and self.api_key.strip()
-            else {"headers": {}}
-        )
 
-        # Build proxy/SSL configuration with environment variable fallback
+        if self.oauth_client_id and self.oauth_client_secret:
+            from pyatlan.client.oauth import OAuthTokenManager
+
+            self._oauth_token_manager = OAuthTokenManager(
+                base_url=str(self.base_url),
+                client_id=self.oauth_client_id,
+                client_secret=self.oauth_client_secret,
+            )
+            self._request_params = {"headers": {}}
+        else:
+            self._request_params = (
+                {"headers": {"authorization": f"Bearer {self.api_key}"}}
+                if self.api_key and self.api_key.strip()
+                else {"headers": {}}
+            )
+
         transport_kwargs = self._build_transport_proxy_config(data)
-        # Configure httpx client with custom transport that supports retry and proxy
-        # Note: We pass proxy/SSL config to the transport, not the client,
-        # so that retry logic properly respects these settings
         self._session = httpx.Client(
             transport=PyatlanSyncTransport(retry=self.retry, **transport_kwargs),
             headers={
@@ -691,7 +701,7 @@ def _call_api_internal(
                     # Retry with impersonation (if _user_id is present)
                     # on authentication failure (token may have expired)
                     if (
-                        self._user_id
+                        (self._user_id or self._oauth_token_manager)
                         and not self._401_has_retried.get()
                         and response.status_code
                         == ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
@@ -813,15 +823,15 @@ def _create_params(
         self, api: API, query_params, request_obj, exclude_unset: bool = True
     ):
         params = copy.deepcopy(self._request_params)
+        if self._oauth_token_manager:
+            token = self._oauth_token_manager.get_token()
+            params["headers"]["authorization"] = f"Bearer {token}"
         params["headers"]["Accept"] = api.consumes
         params["headers"]["content-type"] = api.produces
         if query_params is not None:
             params["params"] = query_params
         if request_obj is not None:
             if isinstance(request_obj, AtlanObject):
-                # Always use AtlanRequest, which accepts a Pydantic model instance and the client
-                # Behind the scenes, it handles retranslation tasks—such as converting
-                # human-readable Atlan tag names back into hashed IDs as required by the backend
                 params["data"] = AtlanRequest(instance=request_obj, client=self).json()
             elif api.consumes == APPLICATION_ENCODED_FORM:
                 params["data"] = request_obj
@@ -838,14 +848,21 @@ def _handle_401_token_refresh(
         download_file_path=None,
         text_response=False,
     ):
-        """
-        Handles token refresh and retries the API request upon a 401 Unauthorized response.
-        1. Impersonates the user (if a user ID is available) to fetch a new token.
-        2. Updates the authorization header with the refreshed token.
-        3. Retries the API request with the new token.
+        if self._oauth_token_manager:
+            self._oauth_token_manager.invalidate_token()
+            token = self._oauth_token_manager.get_token()
+            params["headers"]["authorization"] = f"Bearer {token}"
+            self._401_has_retried.set(True)
+            LOGGER.debug("Successfully refreshed OAuth token after 401.")
+            return self._call_api_internal(
+                api,
+                path,
+                params,
+                binary_data=binary_data,
+                download_file_path=download_file_path,
+                text_response=text_response,
+            )
 
-        returns: HTTP response received after retrying the request with the refreshed token
-        """
         try:
             new_token = self.impersonate.user(user_id=self._user_id)
         except Exception as e:
@@ -861,11 +878,6 @@ def _handle_401_token_refresh(
         self._request_params["headers"]["authorization"] = f"Bearer {self.api_key}"
         LOGGER.debug("Successfully completed 401 automatic token refresh.")
 
-        # Added a retry loop to ensure a token is active before retrying original request
-        # This helps ensure that when we fetch typedefs using the new token,
-        # the backend has fully recognized the token as valid.
-        # Without this delay, we occasionally get an empty response `[]` from the API,
-        # likely because the backend hasn’t fully propagated token validity yet.
         import time
 
         retry_count = 1
@@ -879,10 +891,9 @@ def _handle_401_token_refresh(
                     "Retrying to get typedefs (to ensure token is active) after token refresh failed: %s",
                     e,
                 )
-            time.sleep(retry_count)  # Linear backoff
+            time.sleep(retry_count)
             retry_count += 1
 
-        # Retry the API call with the new token
         return self._call_api_internal(
             api,
             path,

pyatlan/client/oauth.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright 2025 Atlan Pte. Ltd.
+import threading
+import time
+from typing import Optional
+
+import httpx
+from authlib.oauth2.rfc6749 import OAuth2Token
+
+
+class OAuthTokenManager:
+    def __init__(
+        self,
+        base_url: str,
+        client_id: str,
+        client_secret: str,
+        http_client: Optional[httpx.Client] = None,
+    ):
+        self.base_url = base_url.rstrip("/")
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.token_url = f"{self.base_url}/api/service/oauth-clients/token"
+        self._lock = threading.Lock()
+        self._http_client = http_client or httpx.Client(timeout=30.0)
+        self._token: Optional[OAuth2Token] = None
+        self._owns_client = http_client is None
+
+    def get_token(self) -> str:
+        with self._lock:
+            if self._token and not self._token.is_expired():
+                return self._token["access_token"]
+
+            response = self._http_client.post(
+                self.token_url,
+                json={
+                    "clientId": self.client_id,
+                    "clientSecret": self.client_secret,
+                },
+                headers={"Content-Type": "application/json"},
+            )
+            response.raise_for_status()
+
+            data = response.json()
+            access_token = data.get("accessToken") or data.get("access_token")
+            expires_in = data.get("expiresIn") or data.get("expires_in", 600)
+            print(data.get("expiresIn"))
+            self._token = OAuth2Token(
+                {
+                    "access_token": access_token,
+                    "token_type": data.get("tokenType")
+                    or data.get("token_type", "Bearer"),
+                    "expires_in": expires_in,
+                    "expires_at": int(time.time()) + expires_in,
+                }
+            )
+
+            return access_token
+
+    def invalidate_token(self):
+        with self._lock:
+            self._token = None
+
+    def close(self):
+        if self._owns_client:
+            self._http_client.close()

pyproject.toml

@@ -37,6 +37,7 @@ dependencies = [
     "PyYAML~=6.0.3",
     "httpx~=0.28.1",
     "httpx-retries~=0.4.5",
+    "authlib~=1.3.0",
 ]
 
 [project.urls]