feat(pkg): update set_package_headers to accept explicit header values

Accept agent, workflow_id, app_name, and run_id as explicit params,
defaulting to their respective env vars when not provided. The agent
param now defaults to "workflow" instead of reading from X_ATLAN_AGENT.

Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

Author: vaibhavatlan

pyatlan/pkg/utils.py

@@ -191,26 +191,38 @@ def set_package_ops(run_time_config: RuntimeConfig) -> AtlanClient:
     return client
 
 
-def set_package_headers(client: ClientType) -> ClientType:
+def set_package_headers(
+    client: ClientType,
+    agent: str = "workflow",
+    workflow_id: Optional[str] = None,
+    app_name: Optional[str] = None,
+    run_id: Optional[str] = None,
+) -> ClientType:
     """
-    Configure the AtlanClient or AsyncAtlanClient with package headers from environment variables.
+    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)
     :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")
-    ):
+    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": 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", ""
-            ),
+            "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)
     return client

pyatlan_v9/pkg/utils.py

@@ -68,25 +68,38 @@ def get_client(
     return client
 
 
-def set_package_headers(client: AtlanClient) -> AtlanClient:
+def set_package_headers(
+    client: AtlanClient,
+    agent: str = "workflow",
+    workflow_id: Optional[str] = None,
+    app_name: Optional[str] = None,
+    run_id: Optional[str] = None,
+) -> AtlanClient:
     """
-    Configure the AtlanClient with package headers from environment variables.
+    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)
     :returns: updated client instance
     """
-    if (agent := os.environ.get("X_ATLAN_AGENT")) and (
-        agent_id := os.environ.get("X_ATLAN_AGENT_ID")
-    ):
+    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": 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", ""
-            ),
+            "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)
     return client

tests/integration/custom_package_test.py

@@ -100,7 +100,7 @@ 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-workflow-id": "workflow_id_value",
@@ -109,6 +109,50 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
     assert updated_client == mock_client
 
 
+def test_set_package_headers_explicit_values(client: AtlanClient):
+    mock_client = MagicMock(spec=client)
+    updated_client = set_package_headers(
+        mock_client,
+        agent="custom-agent",
+        workflow_id="wf-123",
+        app_name="my-app",
+        run_id="run-456",
+    )
+    expected_headers = {
+        "x-atlan-agent": "custom-agent",
+        "x-atlan-agent-id": "wf-123",
+        "x-atlan-agent-package-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)
+    updated_client = set_package_headers(
+        mock_client,
+        workflow_id="override-wf-id",
+        app_name="override-app",
+    )
+    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",
+    }
+    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,

tests_v9/integration/custom_package_test.py

@@ -100,7 +100,7 @@ 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-workflow-id": "workflow_id_value",
@@ -109,6 +109,50 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
     assert updated_client == mock_client
 
 
+def test_set_package_headers_explicit_values(client: AtlanClient):
+    mock_client = MagicMock(spec=client)
+    updated_client = set_package_headers(
+        mock_client,
+        agent="custom-agent",
+        workflow_id="wf-123",
+        app_name="my-app",
+        run_id="run-456",
+    )
+    expected_headers = {
+        "x-atlan-agent": "custom-agent",
+        "x-atlan-agent-id": "wf-123",
+        "x-atlan-agent-package-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)
+    updated_client = set_package_headers(
+        mock_client,
+        workflow_id="override-wf-id",
+        app_name="override-app",
+    )
+    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",
+    }
+    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,