Merge pull request #707 from atlanhq/APP-8817

APP-8817: Fixed `get_client` bug and bumped to release `8.0.1`

Author: Aryamanz29

.github/workflows/pyatlan-wolfi-base.yaml

@@ -26,6 +26,12 @@ on:
         description: 'Pyatlan version (leave empty to use version.txt ie: latest)'
         required: false
         type: string
+      pyatlan_branch:
+        description: 'Pyatlan git branch (overrides version - installs from git://github.com/atlanhq/atlan-python.git@branch)'
+        required: false
+        type: string
+  release:
+    types: [published]
 
 permissions:
   contents: read
@@ -54,26 +60,54 @@ jobs:
       - name: Set build parameters
         id: set-params
         run: |
-          # Set Python version (default to 3.13 if empty)
-          if [ -z "${{ github.event.inputs.python_version }}" ]; then
+          # Check if triggered by release or manual dispatch
+          if [ "${{ github.event_name }}" = "release" ]; then
+            # Release trigger: use latest Python and SDK version, release build type
+            BUILD_TYPE="release"
             PYTHON_VERSION="3.13"
-          else
-            PYTHON_VERSION="${{ github.event.inputs.python_version }}"
-          fi
-          echo "PYTHON_VERSION=$PYTHON_VERSION" >> $GITHUB_ENV
-          
-          # Set Pyatlan version (default to version.txt if empty)
-          if [ -z "${{ github.event.inputs.pyatlan_version }}" ]; then
             PYATLAN_VERSION=$(cat pyatlan/version.txt)
-            echo "Using pyatlan version from version.txt: $PYATLAN_VERSION"
+            PYATLAN_BRANCH=""
+            INSTALL_FROM_GIT="false"
+            echo "Release trigger detected - using latest versions"
           else
-            PYATLAN_VERSION="${{ github.event.inputs.pyatlan_version }}"
-            echo "Using specified pyatlan version: $PYATLAN_VERSION"
+            # Manual dispatch: use provided inputs with defaults
+            BUILD_TYPE="${{ github.event.inputs.build_type }}"
+            
+            # Set Python version (default to 3.13 if empty)
+            if [ -z "${{ github.event.inputs.python_version }}" ]; then
+              PYTHON_VERSION="3.13"
+            else
+              PYTHON_VERSION="${{ github.event.inputs.python_version }}"
+            fi
+            
+            # Check if branch is provided (overrides version)
+            if [ -n "${{ github.event.inputs.pyatlan_branch }}" ]; then
+              PYATLAN_BRANCH="${{ github.event.inputs.pyatlan_branch }}"
+              PYATLAN_VERSION="branch-${PYATLAN_BRANCH}"
+              INSTALL_FROM_GIT="true"
+              echo "Using pyatlan branch: $PYATLAN_BRANCH"
+            else
+              # Set Pyatlan version (default to version.txt if empty)
+              if [ -z "${{ github.event.inputs.pyatlan_version }}" ]; then
+                PYATLAN_VERSION=$(cat pyatlan/version.txt)
+                echo "Using pyatlan version from version.txt: $PYATLAN_VERSION"
+              else
+                PYATLAN_VERSION="${{ github.event.inputs.pyatlan_version }}"
+                echo "Using specified pyatlan version: $PYATLAN_VERSION"
+              fi
+              PYATLAN_BRANCH=""
+              INSTALL_FROM_GIT="false"
+            fi
           fi
+          
+          echo "BUILD_TYPE=$BUILD_TYPE" >> $GITHUB_ENV
+          echo "PYTHON_VERSION=$PYTHON_VERSION" >> $GITHUB_ENV
           echo "PYATLAN_VERSION=$PYATLAN_VERSION" >> $GITHUB_ENV
+          echo "PYATLAN_BRANCH=$PYATLAN_BRANCH" >> $GITHUB_ENV
+          echo "INSTALL_FROM_GIT=$INSTALL_FROM_GIT" >> $GITHUB_ENV
           
           # Set platforms based on build type
-          if [ "${{ github.event.inputs.build_type }}" = "dev" ]; then
+          if [ "$BUILD_TYPE" = "dev" ]; then
             PLATFORMS="linux/amd64"
           else
             PLATFORMS="linux/amd64,linux/arm64"
@@ -85,16 +119,23 @@ jobs:
           echo "COMMIT_HASH=$COMMIT_HASH" >> $GITHUB_ENV
           
           echo "Build parameters:"
-          echo "  - Build Type: ${{ github.event.inputs.build_type }}"
+          echo "  - Trigger: ${{ github.event_name }}"
+          echo "  - Build Type: $BUILD_TYPE"
           echo "  - Python Version: $PYTHON_VERSION"
           echo "  - Pyatlan Version: $PYATLAN_VERSION"
+          if [ "$INSTALL_FROM_GIT" = "true" ]; then
+            echo "  - Pyatlan Branch: $PYATLAN_BRANCH"
+            echo "  - Install Method: Git (development build)"
+          else
+            echo "  - Install Method: PyPI (stable release)"
+          fi
           echo "  - Platforms: $PLATFORMS"
           echo "  - Commit Hash: $COMMIT_HASH"
 
       - name: Generate image tags
         id: generate-tags
         run: |
-          BUILD_TYPE="${{ github.event.inputs.build_type }}"
+          BUILD_TYPE="${{ env.BUILD_TYPE }}"
           PYTHON_VERSION="${{ env.PYTHON_VERSION }}"
           PYATLAN_VERSION="${{ env.PYATLAN_VERSION }}"
           COMMIT_HASH="${{ env.COMMIT_HASH }}"
@@ -112,7 +153,7 @@ jobs:
           echo "Generated image tag: ghcr.io/atlanhq/pyatlan-wolfi-base:$IMAGE_TAG"
 
       - name: Wait for PyPI availability
-        if: github.event.inputs.pyatlan_version != ''
+        if: env.INSTALL_FROM_GIT == 'false' && (github.event.inputs.pyatlan_version != '' || github.event_name == 'release')
         uses: nick-fields/retry@v3
         with:
           max_attempts: 5
@@ -139,6 +180,8 @@ jobs:
           build-args: |
             PYTHON_VERSION=${{ env.PYTHON_VERSION }}
             PYATLAN_VERSION=${{ env.PYATLAN_VERSION }}
+            PYATLAN_BRANCH=${{ env.PYATLAN_BRANCH }}
+            INSTALL_FROM_GIT=${{ env.INSTALL_FROM_GIT }}
           cache-from: type=gha
           cache-to: type=gha,mode=max
 
@@ -148,11 +191,18 @@ jobs:
           echo "" >> $GITHUB_STEP_SUMMARY
           echo "### Image Details:" >> $GITHUB_STEP_SUMMARY
           echo "- **Base Image:** Wolfi" >> $GITHUB_STEP_SUMMARY
-          echo "- **Build Type:** ${{ github.event.inputs.build_type }}" >> $GITHUB_STEP_SUMMARY
+          echo "- **Trigger:** ${{ github.event_name }}" >> $GITHUB_STEP_SUMMARY
+          echo "- **Build Type:** ${{ env.BUILD_TYPE }}" >> $GITHUB_STEP_SUMMARY
           echo "- **Python Version:** ${{ env.PYTHON_VERSION }}" >> $GITHUB_STEP_SUMMARY
           echo "- **Pyatlan Version:** ${{ env.PYATLAN_VERSION }}" >> $GITHUB_STEP_SUMMARY
+          if [ "${{ env.INSTALL_FROM_GIT }}" = "true" ]; then
+            echo "- **Pyatlan Branch:** ${{ env.PYATLAN_BRANCH }}" >> $GITHUB_STEP_SUMMARY
+            echo "- **Install Method:** Git (development build)" >> $GITHUB_STEP_SUMMARY
+          else
+            echo "- **Install Method:** PyPI (stable release)" >> $GITHUB_STEP_SUMMARY
+          fi
           echo "- **Platforms:** ${{ env.PLATFORMS }}" >> $GITHUB_STEP_SUMMARY
-          if [ "${{ github.event.inputs.build_type }}" = "dev" ]; then
+          if [ "${{ env.BUILD_TYPE }}" = "dev" ]; then
             echo "- **Commit Hash:** ${{ env.COMMIT_HASH }}" >> $GITHUB_STEP_SUMMARY
           fi
           echo "" >> $GITHUB_STEP_SUMMARY

Dockerfile.wolfi

@@ -3,24 +3,41 @@ FROM cgr.dev/chainguard/wolfi-base
 # Build arguments for configurable versions
 ARG PYTHON_VERSION=3.11
 ARG PYATLAN_VERSION=latest
+ARG PYATLAN_BRANCH=""
+ARG INSTALL_FROM_GIT=false
 
 WORKDIR /app
 
-# Install Python
-RUN apk add python-${PYTHON_VERSION} && \
+# Install Python and git (needed for git-based installations)
+RUN apk add python-${PYTHON_VERSION} git && \
     chown -R nonroot:nonroot /app/
 
 # Install uv
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/
 
-# Install pyatlan from PyPI using uv pip install (as root for --system)
+# Install pyatlan - either from PyPI or git branch
 RUN --mount=type=cache,target=/root/.cache/uv \
-    if [ "$PYATLAN_VERSION" = "latest" ]; then \
+    if [ "$INSTALL_FROM_GIT" = "true" ]; then \
+        echo "Installing pyatlan from git branch: $PYATLAN_BRANCH"; \
+        uv pip install --system "git+https://github.com/atlanhq/atlan-python.git@$PYATLAN_BRANCH"; \
+    elif [ "$PYATLAN_VERSION" = "latest" ]; then \
+        echo "Installing latest pyatlan from PyPI"; \
         uv pip install --system pyatlan; \
     else \
+        echo "Installing pyatlan==$PYATLAN_VERSION from PyPI"; \
         uv pip install --system pyatlan==$PYATLAN_VERSION; \
     fi
 
+# Add build information as labels
+RUN if [ "$INSTALL_FROM_GIT" = "true" ]; then \
+        echo "LABEL pyatlan.source=git" >> /tmp/labels && \
+        echo "LABEL pyatlan.branch=$PYATLAN_BRANCH" >> /tmp/labels; \
+    else \
+        echo "LABEL pyatlan.source=pypi" >> /tmp/labels && \
+        echo "LABEL pyatlan.version=$PYATLAN_VERSION" >> /tmp/labels; \
+    fi && \
+    echo "LABEL python.version=$PYTHON_VERSION" >> /tmp/labels
+
 USER nonroot
 
 # Set environment variables

HISTORY.md

@@ -1,3 +1,25 @@
+## 8.0.1 (August 25, 2025)
+
+### New Features
+
+- Added `get_client_async()` for initializing `AsyncAtlanClient` when using packages.
+- Added optional parameter `set_pkg_headers()` to set package headers on the client (defaults to `False`).
+
+### Bug Fixes
+
+- Fixed `httpcore.LocalProtocolError` exception in `get_client()` that occurred when `ATLAN_API_KEY` environment variable was not configured (commonly encountered during package impersonation with empty `API` key strings).
+
+### QOL Improvements
+
+- Added `pyatlan-wolfi-base` Docker image workflow with enhanced capabilities:
+  - **Automatic release builds**: Workflow now triggers automatically on GitHub releases using latest Python (3.13) and SDK versions
+  - **Git branch support**: Added ability to install SDK from any git branch for development/testing purposes via `pyatlan_branch` parameter
+  - **Smart installation logic**: Automatically chooses between PyPI (stable) or git (development) installation methods
+  - **Enhanced tagging**: Branch builds tagged as `branch-{branchname}-{python}-{commit}` for easy identification
+  - **Build metadata**: Images include labels tracking installation source, version/branch, and Python version
+  - **Conditional PyPI checks**: Skips PyPI availability checks when installing from git branches
+  - **Improved logging**: Shows installation method, branch info, and trigger source in build outputs
+
 ## 8.0.0 (August 20, 2025)
 
 ### New Features

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 and self.api_key.strip()
+            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
+) -> AsyncAtlanClient:
+    """
+    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 (

pyatlan/version.txt

@@ -1 +1 @@
-8.0.0
+8.0.1