APP-8817: Fixed get_client() (for both sync and async client errors)

Author: Aryamanz29

pyatlan/client/aio/client.py

@@ -174,10 +174,7 @@ async def from_token_guid(  # type: ignore[override]
 
         # Step 1: Initialize base client and get Atlan-Argo credentials
         # Note: Using empty api_key as we're bootstrapping authentication
-        client = cls(base_url=final_base_url)
-        # Explicitly set api_key to empty string to avoid
-        # httpx.LocalProtocolError: Illegal header value b'Bearer '
-        client.api_key = ""
+        client = cls(base_url=final_base_url, api_key="")
         client_info = ImpersonateUser.get_client_info(
             client_id=client_id, client_secret=client_secret
         )
@@ -873,6 +870,11 @@ async def _upload_file(self, api, file=None, filename=None):
             self._api_logger(api, path)
         return await self._call_api_internal(api, path, params, binary_data=post_data)
 
+    def update_headers(self, header: dict[str, str]):
+        """Update headers for the async session."""
+        if self._async_session:
+            self._async_session.headers.update(header)
+
     async def aclose(self):
         """Close async resources"""
         if self._async_session:

pyatlan/client/atlan.py

@@ -167,11 +167,11 @@ class Config:
 
     def __init__(self, **data):
         super().__init__(**data)
-        self._request_params = {
-            "headers": {
-                "authorization": f"Bearer {self.api_key}",
-            }
-        }
+        self._request_params = (
+            {"headers": {"authorization": f"Bearer {self.api_key}"}}
+            if self.api_key
+            else {"headers": {}}
+        )
         # Configure httpx client with the provided retry settings
         self._session = httpx.Client(
             transport=RetryTransport(retry=self.retry),
@@ -377,10 +377,7 @@ def from_token_guid(
 
         # Step 1: Initialize base client and get Atlan-Argo credentials
         # Note: Using empty api_key as we're bootstrapping authentication
-        client = AtlanClient(base_url=final_base_url)
-        # Explicitly set api_key to empty string to avoid
-        # httpx.LocalProtocolError: Illegal header value b'Bearer '
-        client.api_key = ""
+        client = AtlanClient(base_url=final_base_url, api_key="")
         client_info = ImpersonateUser.get_client_info(
             client_id=client_id, client_secret=client_secret
         )

pyatlan/pkg/utils.py

@@ -4,15 +4,19 @@
 import logging
 import os
 import sys
-from typing import Any, Dict, List, Mapping, Optional, Sequence, Union
+from typing import Any, Dict, List, Mapping, Optional, Sequence, TypeVar, Union
 
 from pydantic.v1 import parse_obj_as, parse_raw_as
 
+from pyatlan.client.aio import AsyncAtlanClient
 from pyatlan.client.atlan import AtlanClient
 from pyatlan.pkg.models import RuntimeConfig
 
 LOGGER = logging.getLogger(__name__)
 
+# Type variable for client types
+ClientType = TypeVar("ClientType", AtlanClient, AsyncAtlanClient)
+
 # Try to import OpenTelemetry libraries
 try:
     from opentelemetry.exporter.otlp.proto.grpc._log_exporter import (  # type:ignore
@@ -64,13 +68,16 @@ def _is_valid_type(self, value: Any) -> bool:
     OTEL_IMPORTS_AVAILABLE = False
 
 
-def get_client(impersonate_user_id: str) -> AtlanClient:
+def get_client(
+    impersonate_user_id: str, set_pkg_headers: Optional[bool] = False
+) -> AtlanClient:
     """
     Set up the default Atlan client, based on environment variables.
     This will use an API token if found in ATLAN_API_KEY, and will fallback to attempting to impersonate a user if
     ATLAN_API_KEY is empty.
 
     :param impersonate_user_id: unique identifier (GUID) of a user or API token to impersonate
+    :param set_pkg_headers: whether to set package headers on the client (default is False)
     :returns: an initialized client
     """
     base_url = os.environ.get("ATLAN_BASE_URL", "INTERNAL")
@@ -94,6 +101,46 @@ def get_client(impersonate_user_id: str) -> AtlanClient:
     client = AtlanClient(base_url=base_url, api_key=api_key)
     if user_id:
         client._user_id = user_id
+    if set_pkg_headers:
+        client = set_package_headers(client)
+    return client
+
+
+async def get_client_async(
+    impersonate_user_id: str, set_pkg_headers: Optional[bool] = False
+):
+    """
+    Set up the default async Atlan client, based on environment variables.
+    This will use an API token if found in ATLAN_API_KEY, and will fallback to attempting to impersonate a user if
+    ATLAN_API_KEY is empty.
+
+    :param impersonate_user_id: unique identifier (GUID) of a user or API token to impersonate
+    :param set_pkg_headers: whether to set package headers on the client (default is False)
+    :returns: an initialized async client
+    """
+    base_url = os.environ.get("ATLAN_BASE_URL", "INTERNAL")
+    api_token = os.environ.get("ATLAN_API_KEY", "")
+    user_id = os.environ.get("ATLAN_USER_ID", impersonate_user_id)
+
+    if api_token:
+        LOGGER.info("Using provided API token for authentication.")
+        api_key = api_token
+    elif user_id:
+        LOGGER.info("No API token found, attempting to impersonate user: %s", user_id)
+        client = AsyncAtlanClient(base_url=base_url, api_key="")
+        api_key = await client.impersonate.user(user_id=user_id)
+    else:
+        LOGGER.info(
+            "No API token or impersonation user, attempting short-lived escalation."
+        )
+        client = AsyncAtlanClient(base_url=base_url, api_key="")
+        api_key = await client.impersonate.escalate()
+
+    client = AsyncAtlanClient(base_url=base_url, api_key=api_key)
+    if user_id:
+        client._user_id = user_id
+    if set_pkg_headers:
+        client = set_package_headers(client)
     return client
 
 
@@ -110,12 +157,12 @@ def set_package_ops(run_time_config: RuntimeConfig) -> AtlanClient:
     return client
 
 
-def set_package_headers(client: AtlanClient) -> AtlanClient:
+def set_package_headers(client: ClientType) -> ClientType:
     """
-    Configure the AtlanClient with package headers from environment variables.
+    Configure the AtlanClient or AsyncAtlanClient with package headers from environment variables.
 
-    :param client: AtlanClient instance to configure
-    :returns: updated AtlanClient instance.
+    :param client: AtlanClient or AsyncAtlanClient instance to configure
+    :returns: updated client instance of the same type.
     """
 
     if (agent := os.environ.get("X_ATLAN_AGENT")) and (

tests/integration/aio/test_client.py

@@ -5,8 +5,10 @@
 
 import pytest
 import pytest_asyncio
+from httpx import Headers
 from pydantic.v1 import StrictStr
 
+from pyatlan import __version__ as VERSION
 from pyatlan.client.aio.client import AsyncAtlanClient
 from pyatlan.client.atlan import DEFAULT_RETRY
 from pyatlan.client.common.audit import LOGGER as AUDIT_LOGGER
@@ -42,6 +44,8 @@
     Term,
 )
 from pyatlan.model.user import UserMinimalResponse
+from pyatlan.pkg.utils import get_client_async
+from pyatlan.utils import get_python_version
 from tests.integration.aio.utils import (
     async_search_with_retry,
     create_database_async,
@@ -1535,6 +1539,80 @@ async def test_client_401_token_refresh(
     assert client.api_key != expired_api_token
     assert results and results.count >= 1
 
+    # Verify similar results with get_client_async()
+    # Setting ATLAN_API_KEY to empty string to force impersonation
+    monkeypatch.setenv("ATLAN_API_KEY", "")
+    assert expired_token_user_id
+    client = await get_client_async(impersonate_user_id=expired_token_user_id)
+    results = await (
+        FluentSearch()
+        .where(CompoundQuery.active_assets())
+        .where(CompoundQuery.asset_type(AtlasGlossary))
+        .page_size(100)
+        .execute_async(client=client)
+    )
+
+    # Confirm the API key has been updated and results are returned
+    assert client.api_key != expired_api_token
+    assert results and results.count >= 1
+
+    # Verify package headers are set correctly
+    expected_common_headers = Headers(
+        {
+            "User-Agent": f"Atlan-PythonSDK/{VERSION}",
+            "Accept-Encoding": "gzip, deflate",
+            "Accept": "*/*",
+            "Connection": "keep-alive",
+            "x-atlan-agent": "sdk",
+            "x-atlan-agent-id": "python",
+            "x-atlan-client-origin": "product_sdk",
+            "x-atlan-python-version": get_python_version(),
+            "x-atlan-client-type": "async",
+        }
+    )
+
+    # Clear package environment variables to test default headers
+    for var in [
+        "X_ATLAN_AGENT",
+        "X_ATLAN_AGENT_ID",
+        "X_ATLAN_AGENT_PACKAGE_NAME",
+        "X_ATLAN_AGENT_WORKFLOW_ID",
+    ]:
+        monkeypatch.delenv(var, raising=False)
+
+    client = await get_client_async(
+        impersonate_user_id=expired_token_user_id, set_pkg_headers=False
+    )
+    assert client._async_session is not None
+    assert expected_common_headers == client._async_session.headers
+
+    # Set package environment variables to test package headers
+    monkeypatch.setenv("X_ATLAN_AGENT", "agent_value")
+    monkeypatch.setenv("X_ATLAN_AGENT_ID", "agent_id_value")
+    monkeypatch.setenv("X_ATLAN_AGENT_PACKAGE_NAME", "package_name_value")
+    monkeypatch.setenv("X_ATLAN_AGENT_WORKFLOW_ID", "workflow_id_value")
+
+    expected = Headers(
+        {
+            "User-Agent": f"Atlan-PythonSDK/{VERSION}",
+            "Accept-Encoding": "gzip, deflate",
+            "Accept": "*/*",
+            "Connection": "keep-alive",
+            "x-atlan-client-origin": "product_sdk",
+            "x-atlan-python-version": get_python_version(),
+            "x-atlan-client-type": "async",
+            "x-atlan-agent": "agent_value",
+            "x-atlan-agent-id": "agent_id_value",
+            "x-atlan-agent-package-name": "package_name_value",
+            "x-atlan-agent-workflow-id": "workflow_id_value",
+        }
+    )
+    client = await get_client_async(
+        impersonate_user_id=expired_token_user_id, set_pkg_headers=True
+    )
+    assert client._async_session is not None
+    assert expected == client._async_session.headers
+
 
 async def test_client_init_from_token_guid(
     client: AsyncAtlanClient, token: ApiToken, argo_fake_token: ApiToken, monkeypatch

tests/integration/test_client.py

@@ -4,8 +4,10 @@
 from unittest.mock import patch
 
 import pytest
+from httpx import Headers
 from pydantic.v1 import StrictStr
 
+from pyatlan import __version__ as VERSION
 from pyatlan.client.atlan import DEFAULT_RETRY, AtlanClient
 from pyatlan.client.common.audit import LOGGER as AUDIT_LOGGER
 from pyatlan.client.common.search_log import LOGGER as SEARCH_LOG_LOGGER
@@ -46,6 +48,8 @@
     Term,
 )
 from pyatlan.model.user import UserMinimalResponse
+from pyatlan.pkg.utils import get_client
+from pyatlan.utils import get_python_version
 from tests.integration.client import TestId
 from tests.integration.lineage_test import create_database, delete_asset
 from tests.integration.requests_test import create_token, delete_token
@@ -1627,6 +1631,76 @@ def test_client_401_token_refresh(
     assert client.api_key != expired_api_token
     assert results and results.count >= 1
 
+    # Verify similar results with get_client()
+    # Setting ATLAN_API_KEY to empty string to force impersonation
+    monkeypatch.setenv("ATLAN_API_KEY", "")
+    assert expired_token_user_id
+    client = get_client(impersonate_user_id=expired_token_user_id)
+    results = (
+        FluentSearch()
+        .where(CompoundQuery.active_assets())
+        .where(CompoundQuery.asset_type(AtlasGlossary))
+        .page_size(100)
+        .execute(client=client)
+    )
+
+    # Confirm the API key has been updated and results are returned
+    assert client.api_key != expired_api_token
+    assert results and results.count >= 1
+
+    # Verify package headers are set correctly
+    expected_common_headers = Headers(
+        {
+            "User-Agent": f"Atlan-PythonSDK/{VERSION}",
+            "Accept-Encoding": "gzip, deflate",
+            "Accept": "*/*",
+            "Connection": "keep-alive",
+            "x-atlan-agent": "sdk",
+            "x-atlan-agent-id": "python",
+            "x-atlan-client-origin": "product_sdk",
+            "x-atlan-python-version": get_python_version(),
+            "x-atlan-client-type": "sync",
+        }
+    )
+
+    # Clear package environment variables to test default headers
+    for var in [
+        "X_ATLAN_AGENT",
+        "X_ATLAN_AGENT_ID",
+        "X_ATLAN_AGENT_PACKAGE_NAME",
+        "X_ATLAN_AGENT_WORKFLOW_ID",
+    ]:
+        monkeypatch.delenv(var, raising=False)
+
+    client = get_client(
+        impersonate_user_id=expired_token_user_id, set_pkg_headers=False
+    )
+    assert expected_common_headers == client._session.headers
+
+    # Set package environment variables to test package headers
+    monkeypatch.setenv("X_ATLAN_AGENT", "agent_value")
+    monkeypatch.setenv("X_ATLAN_AGENT_ID", "agent_id_value")
+    monkeypatch.setenv("X_ATLAN_AGENT_PACKAGE_NAME", "package_name_value")
+    monkeypatch.setenv("X_ATLAN_AGENT_WORKFLOW_ID", "workflow_id_value")
+
+    expected = Headers(
+        {
+            "User-Agent": f"Atlan-PythonSDK/{VERSION}",
+            "Accept-Encoding": "gzip, deflate",
+            "Accept": "*/*",
+            "Connection": "keep-alive",
+            "x-atlan-client-origin": "product_sdk",
+            "x-atlan-python-version": get_python_version(),
+            "x-atlan-client-type": "sync",
+            "x-atlan-agent": "agent_value",
+            "x-atlan-agent-id": "agent_id_value",
+            "x-atlan-agent-package-name": "package_name_value",
+            "x-atlan-agent-workflow-id": "workflow_id_value",
+        }
+    )
+    client = get_client(impersonate_user_id=expired_token_user_id, set_pkg_headers=True)
+    assert expected == client._session.headers
+
 
 def test_client_init_from_token_guid(
     client: AtlanClient, token: ApiToken, argo_fake_token: ApiToken, monkeypatch