feat(pkg): introduce PackageHeaders model for typed header management

Replace loose kwargs in set_package_headers with a PackageHeaders model
(pydantic BaseModel for v1, msgspec.Struct for v9) for better type safety
and maintainability. Rename header key from x-atlan-agent-package-name
to x-atlan-agent-app-name.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: Aryamanz29

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,38 @@
 # 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-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,45 +218,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,
-    agent: str = "workflow",
-    workflow_id: Optional[str] = None,
-    app_name: Optional[str] = None,
-    run_id: Optional[str] = None,
+    headers: Optional[PackageHeaders] = None,
 ) -> ClientType:
     """
     Configure the AtlanClient or AsyncAtlanClient with package headers.
 
-    Each header value can be passed explicitly; if omitted, falls back to the
-    corresponding environment variable (X_ATLAN_AGENT_ID,
-    X_ATLAN_AGENT_PACKAGE_NAME, X_ATLAN_AGENT_WORKFLOW_ID).
-
     :param client: AtlanClient or AsyncAtlanClient instance to configure
-    :param agent: value for the x-atlan-agent header (default: "workflow")
-    :param workflow_id: value for the x-atlan-agent-id header (default: X_ATLAN_AGENT_ID env var)
-    :param app_name: value for the x-atlan-agent-package-name header
-        (default: X_ATLAN_AGENT_PACKAGE_NAME env var)
-    :param run_id: value for the x-atlan-agent-workflow-id header
-        (default: X_ATLAN_AGENT_WORKFLOW_ID env var)
+    :param headers: PackageHeaders instance; if None, reads from environment variables
     :returns: updated client instance of the same type.
     """
-    resolved_workflow_id = workflow_id or os.environ.get("X_ATLAN_AGENT_ID")
-    if agent and resolved_workflow_id:
-        headers: Dict[str, str] = {
-            "x-atlan-agent": agent,
-            "x-atlan-agent-id": resolved_workflow_id,
-            "x-atlan-agent-package-name": app_name
-            or os.environ.get("X_ATLAN_AGENT_PACKAGE_NAME", ""),
-            "x-atlan-agent-workflow-id": run_id
-            or 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
 
 

pyatlan_v9/pkg/utils.py

@@ -5,11 +5,41 @@
 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-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:
@@ -70,36 +100,16 @@ def get_client(
 
 def set_package_headers(
     client: AtlanClient,
-    agent: str = "workflow",
-    workflow_id: Optional[str] = None,
-    app_name: Optional[str] = None,
-    run_id: Optional[str] = None,
+    headers: Optional[PackageHeaders] = None,
 ) -> AtlanClient:
     """
     Configure the AtlanClient with package headers.
 
-    Each header value can be passed explicitly; if omitted, falls back to the
-    corresponding environment variable (X_ATLAN_AGENT_ID,
-    X_ATLAN_AGENT_PACKAGE_NAME, X_ATLAN_AGENT_WORKFLOW_ID).
-
     :param client: AtlanClient instance to configure
-    :param agent: value for the x-atlan-agent header (default: "workflow")
-    :param workflow_id: value for the x-atlan-agent-id header (default: X_ATLAN_AGENT_ID env var)
-    :param app_name: value for the x-atlan-agent-package-name header
-        (default: X_ATLAN_AGENT_PACKAGE_NAME env var)
-    :param run_id: value for the x-atlan-agent-workflow-id header
-        (default: X_ATLAN_AGENT_WORKFLOW_ID env var)
+    :param headers: PackageHeaders instance; if None, reads from environment variables
     :returns: updated client instance
     """
-    resolved_workflow_id = workflow_id or os.environ.get("X_ATLAN_AGENT_ID")
-    if agent and resolved_workflow_id:
-        headers = {
-            "x-atlan-agent": agent,
-            "x-atlan-agent-id": resolved_workflow_id,
-            "x-atlan-agent-package-name": app_name
-            or os.environ.get("X_ATLAN_AGENT_PACKAGE_NAME", ""),
-            "x-atlan-agent-workflow-id": run_id
-            or 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

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",
@@ -102,7 +101,7 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
     expected_headers = {
         "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)
@@ -111,17 +110,17 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
 
 def test_set_package_headers_explicit_values(client: AtlanClient):
     mock_client = MagicMock(spec=client)
-    updated_client = set_package_headers(
-        mock_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)
@@ -130,16 +129,16 @@ def test_set_package_headers_explicit_values(client: AtlanClient):
 
 def test_set_package_headers_explicit_overrides_env(client: AtlanClient, mock_pkg_env):
     mock_client = MagicMock(spec=client)
-    updated_client = set_package_headers(
-        mock_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-workflow-id": "workflow_id_value",
+        "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

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",
@@ -102,7 +101,7 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
     expected_headers = {
         "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)
@@ -111,17 +110,17 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
 
 def test_set_package_headers_explicit_values(client: AtlanClient):
     mock_client = MagicMock(spec=client)
-    updated_client = set_package_headers(
-        mock_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)
@@ -130,16 +129,16 @@ def test_set_package_headers_explicit_values(client: AtlanClient):
 
 def test_set_package_headers_explicit_overrides_env(client: AtlanClient, mock_pkg_env):
     mock_client = MagicMock(spec=client)
-    updated_client = set_package_headers(
-        mock_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-workflow-id": "workflow_id_value",
+        "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