some more fixes

Author: vaibhavatlan

pyatlan/client/aio/client.py

@@ -132,6 +132,7 @@ 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)
 
         if self.oauth_client_id and self.oauth_client_secret:
@@ -143,8 +144,10 @@ def __init__(self, **kwargs):
                 client_secret=self.oauth_client_secret,
             )
 
+        # Build proxy/SSL configuration (reuse from sync client)
         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={
@@ -441,6 +444,9 @@ 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()
@@ -451,6 +457,7 @@ async def _create_params(
             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:
@@ -691,12 +698,17 @@ 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 or self._async_oauth_token_manager)
                 and not self._401_has_retried.get()
                 and response.status_code
                 == ErrorCode.AUTHENTICATION_PASSTHROUGH.http_error_code
             ):
+                """
+                Async version of token refresh and retry logic.
+                Handles token refresh and retries the API request upon a 401 Unauthorized response.
+                """
                 try:
                     LOGGER.debug("Starting async 401 automatic token refresh.")
                     return await self._handle_401_token_refresh(
@@ -763,6 +775,7 @@ async def _handle_401_token_refresh(
             )
 
         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(
@@ -778,9 +791,11 @@ 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]
                 )
@@ -791,9 +806,10 @@ 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)
+            await asyncio.sleep(retry_count)  # Linear backoff with async sleep
             retry_count += 1
 
+        # Retry the API call with the new token
         return await self._call_api_internal(
             api,
             path,

pyatlan/client/atlan.py

@@ -194,7 +194,6 @@ def __init__(self, **data):
 
         # Build proxy/SSL configuration with environment variable fallback
         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
@@ -861,6 +860,7 @@ def _handle_401_token_refresh(
         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.
+
         returns: HTTP response received after retrying the request with the refreshed token
         """
         if self._oauth_token_manager: