Fixes to the code and improvemnt

Author: vaibhavatlan

pyatlan/client/aio/client.py

@@ -146,10 +146,19 @@ def __init__(self, **kwargs):
                 self._oauth_token_manager.close()
                 self._oauth_token_manager = None
 
+            final_base_url = self.base_url or os.environ.get(
+                "ATLAN_BASE_URL", "INTERNAL"
+            )
+            final_oauth_client_id = self.oauth_client_id or os.environ.get(
+                "ATLAN_OAUTH_CLIENT_ID"
+            )
+            final_oauth_client_secret = self.oauth_client_secret or os.environ.get(
+                "ATLAN_OAUTH_CLIENT_SECRET"
+            )
             self._async_oauth_token_manager = AsyncOAuthTokenManager(
-                base_url=str(self.base_url),
-                client_id=self.oauth_client_id,
-                client_secret=self.oauth_client_secret,
+                base_url=final_base_url,
+                client_id=final_oauth_client_id,
+                client_secret=final_oauth_client_secret,
             )
 
         # Build proxy/SSL configuration (reuse from sync client)

pyatlan/client/aio/oauth.py

@@ -7,6 +7,9 @@
 import httpx
 from authlib.oauth2.rfc6749 import OAuth2Token
 
+from pyatlan.client.constants import GET_OAUTH_CLIENT
+from pyatlan.utils import API
+
 
 class AsyncOAuthTokenManager:
     def __init__(
@@ -16,10 +19,10 @@ def __init__(
         client_secret: str,
         http_client: Optional[httpx.AsyncClient] = None,
     ):
-        self.base_url = base_url.rstrip("/")
+        self.base_url = base_url
         self.client_id = client_id
         self.client_secret = client_secret
-        self.token_url = f"{self.base_url}/api/service/oauth-clients/token"
+        self.token_url = self._create_path(GET_OAUTH_CLIENT)
         self._lock = asyncio.Lock()
         self._http_client = http_client or httpx.AsyncClient(timeout=30.0)
         self._token: Optional[OAuth2Token] = None
@@ -28,7 +31,7 @@ def __init__(
     async def get_token(self) -> str:
         async with self._lock:
             if self._token and not self._token.is_expired():
-                return self._token["access_token"]
+                return str(self._token["access_token"])
 
             response = await self._http_client.post(
                 self.token_url,
@@ -67,6 +70,16 @@ async def invalidate_token(self):
         async with self._lock:
             self._token = None
 
+    def _create_path(self, api: API):
+        from urllib.parse import urljoin
+
+        if self.base_url == "INTERNAL":
+            base_with_prefix = urljoin(api.endpoint.service, api.endpoint.prefix)
+            return urljoin(base_with_prefix, api.path)
+        else:
+            base_with_prefix = urljoin(self.base_url, api.endpoint.prefix)
+            return urljoin(base_with_prefix, api.path)
+
     async def aclose(self):
         if self._owns_client:
             await self._http_client.aclose()

pyatlan/client/atlan.py

@@ -48,6 +48,7 @@
 from pyatlan.client.file import FileClient
 from pyatlan.client.group import GroupClient
 from pyatlan.client.impersonate import ImpersonationClient
+from pyatlan.client.oauth import OAuthTokenManager
 from pyatlan.client.open_lineage import OpenLineageClient
 from pyatlan.client.query import QueryClient
 from pyatlan.client.role import RoleClient
@@ -177,12 +178,21 @@ def __init__(self, **data):
         super().__init__(**data)
 
         if self.oauth_client_id and self.oauth_client_secret and self.api_key is None:
-            from pyatlan.client.oauth import OAuthTokenManager
+            LOGGER.debug("API KEY not provided. Using OAuth flow for authentication")
 
+            final_base_url = self.base_url or os.environ.get(
+                "ATLAN_BASE_URL", "INTERNAL"
+            )
+            final_oauth_client_id = self.oauth_client_id or os.environ.get(
+                "ATLAN_OAUTH_CLIENT_ID"
+            )
+            final_oauth_client_secret = self.oauth_client_secret or os.environ.get(
+                "ATLAN_OAUTH_CLIENT_SECRET"
+            )
             self._oauth_token_manager = OAuthTokenManager(
-                base_url=str(self.base_url),
-                client_id=self.oauth_client_id,
-                client_secret=self.oauth_client_secret,
+                base_url=final_base_url,
+                client_id=final_oauth_client_id,
+                client_secret=final_oauth_client_secret,
             )
             self._request_params = {"headers": {}}
         else:

pyatlan/client/constants.py

@@ -88,6 +88,14 @@
 GET_WHOAMI_USER = API(
     WHOAMI_API, HTTPMethod.GET, HTTPStatus.OK, endpoint=EndPoint.HERACLES
 )
+
+# oauth client authentinatication
+GET_OAUTH_CLIENT = API(
+    "oauth-clients/token",
+    HTTPMethod.POST,
+    HTTPStatus.OK,
+    endpoint=EndPoint.HERACLES,
+)
 # SQL parsing APIs
 PARSE_QUERY = API(
     f"{QUERY_API}/parse", HTTPMethod.POST, HTTPStatus.OK, endpoint=EndPoint.HEKA

pyatlan/client/oauth.py

@@ -7,6 +7,9 @@
 import httpx
 from authlib.oauth2.rfc6749 import OAuth2Token
 
+from pyatlan.client.constants import GET_OAUTH_CLIENT
+from pyatlan.utils import API
+
 
 class OAuthTokenManager:
     def __init__(
@@ -16,10 +19,10 @@ def __init__(
         client_secret: str,
         http_client: Optional[httpx.Client] = None,
     ):
-        self.base_url = base_url.rstrip("/")
+        self.base_url = base_url
         self.client_id = client_id
         self.client_secret = client_secret
-        self.token_url = f"{self.base_url}/api/service/oauth-clients/token"
+        self.token_url = self._create_path(GET_OAUTH_CLIENT)
         self._lock = threading.Lock()
         self._http_client = http_client or httpx.Client(timeout=30.0)
         self._token: Optional[OAuth2Token] = None
@@ -28,7 +31,7 @@ def __init__(
     def get_token(self) -> str:
         with self._lock:
             if self._token and not self._token.is_expired():
-                return self._token["access_token"]
+                return str(self._token["access_token"])
 
             response = self._http_client.post(
                 self.token_url,
@@ -67,6 +70,16 @@ def invalidate_token(self):
         with self._lock:
             self._token = None
 
+    def _create_path(self, api: API):
+        from urllib.parse import urljoin
+
+        if self.base_url == "INTERNAL":
+            base_with_prefix = urljoin(api.endpoint.service, api.endpoint.prefix)
+            return urljoin(base_with_prefix, api.path)
+        else:
+            base_with_prefix = urljoin(self.base_url, api.endpoint.prefix)
+            return urljoin(base_with_prefix, api.path)
+
     def close(self):
         if self._owns_client:
             self._http_client.close()

pyproject.toml

@@ -51,6 +51,7 @@ dev = [
     "mypy~=1.18.0",
     "ruff~=0.14.0",
     "types-setuptools~=80.9.0.20250822",
+    "types-Authlib",
     "pytest~=8.4.2",
     "pytest-vcr~=1.0.2",
     "vcrpy~=7.0.0",