Merge pull request #847 from atlanhq/BLDX-728

Aryamanz29 · web-flow · commit 3969ddb3f21c · 2026-03-11T13:21:05.000+05:30
BLDX-728 |  update set_package_headers to accept explicit header values
diff --git a/pyatlan/pkg/utils.py b/pyatlan/pkg/utils.py
@@ -6,7 +6,7 @@
 import sys
 from typing import Any, Dict, List, Mapping, Optional, Sequence, TypeVar, Union
 
-from pydantic.v1 import parse_obj_as, parse_raw_as
+from pydantic.v1 import BaseModel, parse_obj_as, parse_raw_as
 
 from pyatlan.client.aio import AsyncAtlanClient
 from pyatlan.client.atlan import AtlanClient
@@ -17,6 +17,39 @@
 # Type variable for client types
 ClientType = TypeVar("ClientType", AtlanClient, AsyncAtlanClient)
 
+
+class PackageHeaders(BaseModel):
+    """Typed container for Atlan package HTTP headers."""
+
+    agent: str = "workflow"
+    workflow_id: Optional[str] = None
+    app_name: Optional[str] = None
+    run_id: Optional[str] = None
+
+    class Config:
+        allow_mutation = False
+
+    def to_headers(self) -> Dict[str, str]:
+        return {
+            "x-atlan-agent": self.agent,
+            "x-atlan-agent-id": self.workflow_id or "",
+            "x-atlan-agent-package-name": self.app_name or "",
+            "x-atlan-agent-app-name": self.app_name or "",
+            "x-atlan-agent-workflow-id": self.run_id or "",
+        }
+
+    @classmethod
+    def from_env(cls) -> Optional["PackageHeaders"]:
+        workflow_id = os.environ.get("X_ATLAN_AGENT_ID")
+        if not workflow_id:
+            return None
+        return cls(
+            workflow_id=workflow_id,
+            app_name=os.environ.get("X_ATLAN_AGENT_PACKAGE_NAME", ""),
+            run_id=os.environ.get("X_ATLAN_AGENT_WORKFLOW_ID", ""),
+        )
+
+
 # Try to import OpenTelemetry libraries
 try:
     from opentelemetry.exporter.otlp.proto.grpc._log_exporter import (  # type:ignore
@@ -186,33 +219,33 @@ def set_package_ops(run_time_config: RuntimeConfig) -> AtlanClient:
     :returns: an intialized AtlanClient that should be used for any calls to the SDK
     """
     client = get_client(run_time_config.user_id or "")
-    if run_time_config.agent == "workflow":
-        client = set_package_headers(client)
+    if run_time_config.agent:
+        client = set_package_headers(
+            client,
+            headers=PackageHeaders(
+                agent=run_time_config.agent,
+                workflow_id=run_time_config.agent_id,
+                app_name=run_time_config.agent_pkg,
+                run_id=run_time_config.agent_wfl,
+            ),
+        )
     return client
 
 
-def set_package_headers(client: ClientType) -> ClientType:
+def set_package_headers(
+    client: ClientType,
+    headers: Optional[PackageHeaders] = None,
+) -> ClientType:
     """
-    Configure the AtlanClient or AsyncAtlanClient with package headers from environment variables.
+    Configure the AtlanClient or AsyncAtlanClient with package headers.
 
     :param client: AtlanClient or AsyncAtlanClient instance to configure
+    :param headers: PackageHeaders instance; if None, reads from environment variables
     :returns: updated client instance of the same type.
     """
-
-    if (agent := os.environ.get("X_ATLAN_AGENT")) and (
-        agent_id := os.environ.get("X_ATLAN_AGENT_ID")
-    ):
-        headers: Dict[str, str] = {
-            "x-atlan-agent": agent,
-            "x-atlan-agent-id": agent_id,
-            "x-atlan-agent-package-name": os.environ.get(
-                "X_ATLAN_AGENT_PACKAGE_NAME", ""
-            ),
-            "x-atlan-agent-workflow-id": os.environ.get(
-                "X_ATLAN_AGENT_WORKFLOW_ID", ""
-            ),
-        }
-        client.update_headers(headers)
+    pkg_headers = headers or PackageHeaders.from_env()
+    if pkg_headers and pkg_headers.workflow_id:
+        client.update_headers(pkg_headers.to_headers())
     return client
 
 
diff --git a/pyatlan_v9/pkg/utils.py b/pyatlan_v9/pkg/utils.py
@@ -5,11 +5,42 @@
 import os
 from typing import Optional
 
+import msgspec
+
 from pyatlan_v9.client.atlan import AtlanClient
 
 LOGGER = logging.getLogger(__name__)
 
 
+class PackageHeaders(msgspec.Struct, frozen=True, kw_only=True):
+    """Typed container for Atlan package HTTP headers."""
+
+    agent: str = "workflow"
+    workflow_id: Optional[str] = None
+    app_name: Optional[str] = None
+    run_id: Optional[str] = None
+
+    def to_headers(self) -> dict[str, str]:
+        return {
+            "x-atlan-agent": self.agent,
+            "x-atlan-agent-id": self.workflow_id or "",
+            "x-atlan-agent-package-name": self.app_name or "",
+            "x-atlan-agent-app-name": self.app_name or "",
+            "x-atlan-agent-workflow-id": self.run_id or "",
+        }
+
+    @classmethod
+    def from_env(cls) -> Optional["PackageHeaders"]:
+        workflow_id = os.environ.get("X_ATLAN_AGENT_ID")
+        if not workflow_id:
+            return None
+        return cls(
+            workflow_id=workflow_id,
+            app_name=os.environ.get("X_ATLAN_AGENT_PACKAGE_NAME", ""),
+            run_id=os.environ.get("X_ATLAN_AGENT_WORKFLOW_ID", ""),
+        )
+
+
 def get_client(
     impersonate_user_id: Optional[str] = None, set_pkg_headers: Optional[bool] = False
 ) -> AtlanClient:
@@ -68,25 +99,18 @@ def get_client(
     return client
 
 
-def set_package_headers(client: AtlanClient) -> AtlanClient:
+def set_package_headers(
+    client: AtlanClient,
+    headers: Optional[PackageHeaders] = None,
+) -> AtlanClient:
     """
-    Configure the AtlanClient with package headers from environment variables.
+    Configure the AtlanClient with package headers.
 
     :param client: AtlanClient instance to configure
+    :param headers: PackageHeaders instance; if None, reads from environment variables
     :returns: updated client instance
     """
-    if (agent := os.environ.get("X_ATLAN_AGENT")) and (
-        agent_id := os.environ.get("X_ATLAN_AGENT_ID")
-    ):
-        headers = {
-            "x-atlan-agent": agent,
-            "x-atlan-agent-id": agent_id,
-            "x-atlan-agent-package-name": os.environ.get(
-                "X_ATLAN_AGENT_PACKAGE_NAME", ""
-            ),
-            "x-atlan-agent-workflow-id": os.environ.get(
-                "X_ATLAN_AGENT_WORKFLOW_ID", ""
-            ),
-        }
-        client.update_headers(headers)
+    pkg_headers = headers or PackageHeaders.from_env()
+    if pkg_headers and pkg_headers.workflow_id:
+        client.update_headers(pkg_headers.to_headers())
     return client
diff --git a/tests/integration/custom_package_test.py b/tests/integration/custom_package_test.py
@@ -9,7 +9,7 @@
 from pyatlan.client.impersonate import ImpersonationClient
 from pyatlan.pkg.models import CustomPackage, generate
 from pyatlan.pkg.ui import UIConfig, UIStep
-from pyatlan.pkg.utils import get_client, set_package_headers
+from pyatlan.pkg.utils import PackageHeaders, get_client, set_package_headers
 from pyatlan.pkg.widgets import TextInput
 
 
@@ -18,7 +18,6 @@ def mock_pkg_env():
     with patch.dict(
         os.environ,
         {
-            "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",
@@ -100,15 +99,62 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
     mock_client = MagicMock(spec=client)
     updated_client = set_package_headers(mock_client)
     expected_headers = {
-        "x-atlan-agent": "agent_value",
+        "x-atlan-agent": "workflow",
         "x-atlan-agent-id": "agent_id_value",
         "x-atlan-agent-package-name": "package_name_value",
+        "x-atlan-agent-app-name": "package_name_value",
         "x-atlan-agent-workflow-id": "workflow_id_value",
     }
     mock_client.update_headers.assert_called_once_with(expected_headers)
     assert updated_client == mock_client
 
 
+def test_set_package_headers_explicit_values(client: AtlanClient):
+    mock_client = MagicMock(spec=client)
+    headers = PackageHeaders(
+        agent="custom-agent",
+        workflow_id="wf-123",
+        app_name="my-app",
+        run_id="run-456",
+    )
+    updated_client = set_package_headers(mock_client, headers=headers)
+    expected_headers = {
+        "x-atlan-agent": "custom-agent",
+        "x-atlan-agent-id": "wf-123",
+        "x-atlan-agent-package-name": "my-app",
+        "x-atlan-agent-app-name": "my-app",
+        "x-atlan-agent-workflow-id": "run-456",
+    }
+    mock_client.update_headers.assert_called_once_with(expected_headers)
+    assert updated_client == mock_client
+
+
+def test_set_package_headers_explicit_overrides_env(client: AtlanClient, mock_pkg_env):
+    mock_client = MagicMock(spec=client)
+    headers = PackageHeaders(
+        workflow_id="override-wf-id",
+        app_name="override-app",
+    )
+    updated_client = set_package_headers(mock_client, headers=headers)
+    expected_headers = {
+        "x-atlan-agent": "workflow",
+        "x-atlan-agent-id": "override-wf-id",
+        "x-atlan-agent-package-name": "override-app",
+        "x-atlan-agent-app-name": "override-app",
+        "x-atlan-agent-workflow-id": "",
+    }
+    mock_client.update_headers.assert_called_once_with(expected_headers)
+    assert updated_client == mock_client
+
+
+def test_set_package_headers_no_workflow_id(client: AtlanClient):
+    mock_client = MagicMock(spec=client)
+    with patch.dict(os.environ, {}, clear=True):
+        updated_client = set_package_headers(mock_client)
+    mock_client.update_headers.assert_not_called()
+    assert updated_client == mock_client
+
+
 @patch.object(ImpersonationClient, "user", return_value="some-api-key")
 def test_get_client_user_id_handling(
     mock_impersonate_client,
diff --git a/tests_v9/integration/custom_package_test.py b/tests_v9/integration/custom_package_test.py
@@ -10,15 +10,14 @@
 from pyatlan.pkg.widgets import TextInput
 from pyatlan_v9.client.atlan import AtlanClient
 from pyatlan_v9.client.impersonate import V9ImpersonationClient
-from pyatlan_v9.pkg.utils import get_client, set_package_headers
+from pyatlan_v9.pkg.utils import PackageHeaders, get_client, set_package_headers
 
 
 @pytest.fixture
 def mock_pkg_env():
     with patch.dict(
         os.environ,
         {
-            "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",
@@ -100,15 +99,62 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
     mock_client = MagicMock(spec=client)
     updated_client = set_package_headers(mock_client)
     expected_headers = {
-        "x-atlan-agent": "agent_value",
+        "x-atlan-agent": "workflow",
         "x-atlan-agent-id": "agent_id_value",
         "x-atlan-agent-package-name": "package_name_value",
+        "x-atlan-agent-app-name": "package_name_value",
         "x-atlan-agent-workflow-id": "workflow_id_value",
     }
     mock_client.update_headers.assert_called_once_with(expected_headers)
     assert updated_client == mock_client
 
 
+def test_set_package_headers_explicit_values(client: AtlanClient):
+    mock_client = MagicMock(spec=client)
+    headers = PackageHeaders(
+        agent="custom-agent",
+        workflow_id="wf-123",
+        app_name="my-app",
+        run_id="run-456",
+    )
+    updated_client = set_package_headers(mock_client, headers=headers)
+    expected_headers = {
+        "x-atlan-agent": "custom-agent",
+        "x-atlan-agent-id": "wf-123",
+        "x-atlan-agent-package-name": "my-app",
+        "x-atlan-agent-app-name": "my-app",
+        "x-atlan-agent-workflow-id": "run-456",
+    }
+    mock_client.update_headers.assert_called_once_with(expected_headers)
+    assert updated_client == mock_client
+
+
+def test_set_package_headers_explicit_overrides_env(client: AtlanClient, mock_pkg_env):
+    mock_client = MagicMock(spec=client)
+    headers = PackageHeaders(
+        workflow_id="override-wf-id",
+        app_name="override-app",
+    )
+    updated_client = set_package_headers(mock_client, headers=headers)
+    expected_headers = {
+        "x-atlan-agent": "workflow",
+        "x-atlan-agent-id": "override-wf-id",
+        "x-atlan-agent-package-name": "override-app",
+        "x-atlan-agent-app-name": "override-app",
+        "x-atlan-agent-workflow-id": "",
+    }
+    mock_client.update_headers.assert_called_once_with(expected_headers)
+    assert updated_client == mock_client
+
+
+def test_set_package_headers_no_workflow_id(client: AtlanClient):
+    mock_client = MagicMock(spec=client)
+    with patch.dict(os.environ, {}, clear=True):
+        updated_client = set_package_headers(mock_client)
+    mock_client.update_headers.assert_not_called()
+    assert updated_client == mock_client
+
+
 @patch.object(V9ImpersonationClient, "user", return_value="some-api-key")
 def test_get_client_user_id_handling(
     mock_impersonate_client,
