Skip to content

Commit af34cbe

Browse files
vaibhavatlanclaude
andcommitted
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>
1 parent 70a085d commit af34cbe

4 files changed

Lines changed: 140 additions & 27 deletions

File tree

pyatlan/pkg/utils.py

Lines changed: 25 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -191,26 +191,38 @@ def set_package_ops(run_time_config: RuntimeConfig) -> AtlanClient:
191191
return client
192192

193193

194-
def set_package_headers(client: ClientType) -> ClientType:
194+
def set_package_headers(
195+
client: ClientType,
196+
agent: str = "workflow",
197+
workflow_id: Optional[str] = None,
198+
app_name: Optional[str] = None,
199+
run_id: Optional[str] = None,
200+
) -> ClientType:
195201
"""
196-
Configure the AtlanClient or AsyncAtlanClient with package headers from environment variables.
202+
Configure the AtlanClient or AsyncAtlanClient with package headers.
203+
204+
Each header value can be passed explicitly; if omitted, falls back to the
205+
corresponding environment variable (X_ATLAN_AGENT_ID,
206+
X_ATLAN_AGENT_PACKAGE_NAME, X_ATLAN_AGENT_WORKFLOW_ID).
197207
198208
:param client: AtlanClient or AsyncAtlanClient instance to configure
209+
:param agent: value for the x-atlan-agent header (default: "workflow")
210+
:param workflow_id: value for the x-atlan-agent-id header (default: X_ATLAN_AGENT_ID env var)
211+
:param app_name: value for the x-atlan-agent-package-name header
212+
(default: X_ATLAN_AGENT_PACKAGE_NAME env var)
213+
:param run_id: value for the x-atlan-agent-workflow-id header
214+
(default: X_ATLAN_AGENT_WORKFLOW_ID env var)
199215
:returns: updated client instance of the same type.
200216
"""
201-
202-
if (agent := os.environ.get("X_ATLAN_AGENT")) and (
203-
agent_id := os.environ.get("X_ATLAN_AGENT_ID")
204-
):
217+
resolved_workflow_id = workflow_id or os.environ.get("X_ATLAN_AGENT_ID")
218+
if agent and resolved_workflow_id:
205219
headers: Dict[str, str] = {
206220
"x-atlan-agent": agent,
207-
"x-atlan-agent-id": agent_id,
208-
"x-atlan-agent-package-name": os.environ.get(
209-
"X_ATLAN_AGENT_PACKAGE_NAME", ""
210-
),
211-
"x-atlan-agent-workflow-id": os.environ.get(
212-
"X_ATLAN_AGENT_WORKFLOW_ID", ""
213-
),
221+
"x-atlan-agent-id": resolved_workflow_id,
222+
"x-atlan-agent-package-name": app_name
223+
or os.environ.get("X_ATLAN_AGENT_PACKAGE_NAME", ""),
224+
"x-atlan-agent-workflow-id": run_id
225+
or os.environ.get("X_ATLAN_AGENT_WORKFLOW_ID", ""),
214226
}
215227
client.update_headers(headers)
216228
return client

pyatlan_v9/pkg/utils.py

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -68,25 +68,38 @@ def get_client(
6868
return client
6969

7070

71-
def set_package_headers(client: AtlanClient) -> AtlanClient:
71+
def set_package_headers(
72+
client: AtlanClient,
73+
agent: str = "workflow",
74+
workflow_id: Optional[str] = None,
75+
app_name: Optional[str] = None,
76+
run_id: Optional[str] = None,
77+
) -> AtlanClient:
7278
"""
73-
Configure the AtlanClient with package headers from environment variables.
79+
Configure the AtlanClient with package headers.
80+
81+
Each header value can be passed explicitly; if omitted, falls back to the
82+
corresponding environment variable (X_ATLAN_AGENT_ID,
83+
X_ATLAN_AGENT_PACKAGE_NAME, X_ATLAN_AGENT_WORKFLOW_ID).
7484
7585
:param client: AtlanClient instance to configure
86+
:param agent: value for the x-atlan-agent header (default: "workflow")
87+
:param workflow_id: value for the x-atlan-agent-id header (default: X_ATLAN_AGENT_ID env var)
88+
:param app_name: value for the x-atlan-agent-package-name header
89+
(default: X_ATLAN_AGENT_PACKAGE_NAME env var)
90+
:param run_id: value for the x-atlan-agent-workflow-id header
91+
(default: X_ATLAN_AGENT_WORKFLOW_ID env var)
7692
:returns: updated client instance
7793
"""
78-
if (agent := os.environ.get("X_ATLAN_AGENT")) and (
79-
agent_id := os.environ.get("X_ATLAN_AGENT_ID")
80-
):
94+
resolved_workflow_id = workflow_id or os.environ.get("X_ATLAN_AGENT_ID")
95+
if agent and resolved_workflow_id:
8196
headers = {
8297
"x-atlan-agent": agent,
83-
"x-atlan-agent-id": agent_id,
84-
"x-atlan-agent-package-name": os.environ.get(
85-
"X_ATLAN_AGENT_PACKAGE_NAME", ""
86-
),
87-
"x-atlan-agent-workflow-id": os.environ.get(
88-
"X_ATLAN_AGENT_WORKFLOW_ID", ""
89-
),
98+
"x-atlan-agent-id": resolved_workflow_id,
99+
"x-atlan-agent-package-name": app_name
100+
or os.environ.get("X_ATLAN_AGENT_PACKAGE_NAME", ""),
101+
"x-atlan-agent-workflow-id": run_id
102+
or os.environ.get("X_ATLAN_AGENT_WORKFLOW_ID", ""),
90103
}
91104
client.update_headers(headers)
92105
return client

tests/integration/custom_package_test.py

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
100100
mock_client = MagicMock(spec=client)
101101
updated_client = set_package_headers(mock_client)
102102
expected_headers = {
103-
"x-atlan-agent": "agent_value",
103+
"x-atlan-agent": "workflow",
104104
"x-atlan-agent-id": "agent_id_value",
105105
"x-atlan-agent-package-name": "package_name_value",
106106
"x-atlan-agent-workflow-id": "workflow_id_value",
@@ -109,6 +109,50 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
109109
assert updated_client == mock_client
110110

111111

112+
def test_set_package_headers_explicit_values(client: AtlanClient):
113+
mock_client = MagicMock(spec=client)
114+
updated_client = set_package_headers(
115+
mock_client,
116+
agent="custom-agent",
117+
workflow_id="wf-123",
118+
app_name="my-app",
119+
run_id="run-456",
120+
)
121+
expected_headers = {
122+
"x-atlan-agent": "custom-agent",
123+
"x-atlan-agent-id": "wf-123",
124+
"x-atlan-agent-package-name": "my-app",
125+
"x-atlan-agent-workflow-id": "run-456",
126+
}
127+
mock_client.update_headers.assert_called_once_with(expected_headers)
128+
assert updated_client == mock_client
129+
130+
131+
def test_set_package_headers_explicit_overrides_env(client: AtlanClient, mock_pkg_env):
132+
mock_client = MagicMock(spec=client)
133+
updated_client = set_package_headers(
134+
mock_client,
135+
workflow_id="override-wf-id",
136+
app_name="override-app",
137+
)
138+
expected_headers = {
139+
"x-atlan-agent": "workflow",
140+
"x-atlan-agent-id": "override-wf-id",
141+
"x-atlan-agent-package-name": "override-app",
142+
"x-atlan-agent-workflow-id": "workflow_id_value",
143+
}
144+
mock_client.update_headers.assert_called_once_with(expected_headers)
145+
assert updated_client == mock_client
146+
147+
148+
def test_set_package_headers_no_workflow_id(client: AtlanClient):
149+
mock_client = MagicMock(spec=client)
150+
with patch.dict(os.environ, {}, clear=True):
151+
updated_client = set_package_headers(mock_client)
152+
mock_client.update_headers.assert_not_called()
153+
assert updated_client == mock_client
154+
155+
112156
@patch.object(ImpersonationClient, "user", return_value="some-api-key")
113157
def test_get_client_user_id_handling(
114158
mock_impersonate_client,

tests_v9/integration/custom_package_test.py

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
100100
mock_client = MagicMock(spec=client)
101101
updated_client = set_package_headers(mock_client)
102102
expected_headers = {
103-
"x-atlan-agent": "agent_value",
103+
"x-atlan-agent": "workflow",
104104
"x-atlan-agent-id": "agent_id_value",
105105
"x-atlan-agent-package-name": "package_name_value",
106106
"x-atlan-agent-workflow-id": "workflow_id_value",
@@ -109,6 +109,50 @@ def test_set_package_headers(client: AtlanClient, mock_pkg_env):
109109
assert updated_client == mock_client
110110

111111

112+
def test_set_package_headers_explicit_values(client: AtlanClient):
113+
mock_client = MagicMock(spec=client)
114+
updated_client = set_package_headers(
115+
mock_client,
116+
agent="custom-agent",
117+
workflow_id="wf-123",
118+
app_name="my-app",
119+
run_id="run-456",
120+
)
121+
expected_headers = {
122+
"x-atlan-agent": "custom-agent",
123+
"x-atlan-agent-id": "wf-123",
124+
"x-atlan-agent-package-name": "my-app",
125+
"x-atlan-agent-workflow-id": "run-456",
126+
}
127+
mock_client.update_headers.assert_called_once_with(expected_headers)
128+
assert updated_client == mock_client
129+
130+
131+
def test_set_package_headers_explicit_overrides_env(client: AtlanClient, mock_pkg_env):
132+
mock_client = MagicMock(spec=client)
133+
updated_client = set_package_headers(
134+
mock_client,
135+
workflow_id="override-wf-id",
136+
app_name="override-app",
137+
)
138+
expected_headers = {
139+
"x-atlan-agent": "workflow",
140+
"x-atlan-agent-id": "override-wf-id",
141+
"x-atlan-agent-package-name": "override-app",
142+
"x-atlan-agent-workflow-id": "workflow_id_value",
143+
}
144+
mock_client.update_headers.assert_called_once_with(expected_headers)
145+
assert updated_client == mock_client
146+
147+
148+
def test_set_package_headers_no_workflow_id(client: AtlanClient):
149+
mock_client = MagicMock(spec=client)
150+
with patch.dict(os.environ, {}, clear=True):
151+
updated_client = set_package_headers(mock_client)
152+
mock_client.update_headers.assert_not_called()
153+
assert updated_client == mock_client
154+
155+
112156
@patch.object(V9ImpersonationClient, "user", return_value="some-api-key")
113157
def test_get_client_user_id_handling(
114158
mock_impersonate_client,

0 commit comments

Comments
 (0)