You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by vi...@apache.org on 2026/04/27 14:28:03 UTC
(airflow) branch main updated: Add multi-team dimensions to boto3 user agent string (#65821)
This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f928801b953 Add multi-team dimensions to boto3 user agent string (#65821)
f928801b953 is described below
commit f928801b95342ccae1e1ed046af89111ac7bbe79
Author: Niko Oliveira <on...@amazon.com>
AuthorDate: Mon Apr 27 07:27:53 2026 -0700
Add multi-team dimensions to boto3 user agent string (#65821)
Add MultiTeam and TeamNameKey dimensions to the user agent extra field
in AwsBaseHook. MultiTeam indicates whether the task is running under
multi-team Airflow (True/False). TeamNameKey is a SHA-1 hashed UUID of
the team name, following the same privacy-preserving pattern as
DagRunKey. Both are derived from the AIRFLOW_CTX_TEAM_NAME environment
variable propagated by the execution API.
---
.../airflow/providers/amazon/aws/hooks/base_aws.py | 13 +++++++++++
.../tests/unit/amazon/aws/hooks/test_base_aws.py | 25 +++++++++++++++++++++-
2 files changed, 37 insertions(+), 1 deletion(-)
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py
index 5fc6c4a9a70..59249fac911 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/base_aws.py
@@ -605,12 +605,25 @@ class AwsGenericHook(BaseHook, Generic[BaseAwsConnection]):
return airflow_version
+ @staticmethod
+ def _is_multi_team() -> bool:
+ """Return True if the current task is running under multi-team Airflow."""
+ return bool(os.environ.get("AIRFLOW_CTX_TEAM_NAME"))
+
+ @staticmethod
+ @return_on_error("00000000-0000-0000-0000-000000000000")
+ def _generate_team_name_key() -> str:
+ """Generate a hashed key from the team name, using the same approach as :meth:`_generate_dag_key`."""
+ return generate_uuid(os.environ.get("AIRFLOW_CTX_TEAM_NAME"))
+
def _generate_user_agent_extra_field(self, existing_user_agent_extra: str) -> str:
user_agent_extra_values = [
f"Airflow/{self._get_airflow_version()}",
f"AmPP/{self._get_provider_version()}",
f"Caller/{self._get_caller()}",
f"DagRunKey/{self._generate_dag_key()}",
+ f"MultiTeam/{self._is_multi_team()}",
+ f"TeamNameKey/{self._generate_team_name_key()}",
existing_user_agent_extra or "",
]
return " ".join(user_agent_extra_values).strip()
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py
index d2fa7776efd..e0984f88e24 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_base_aws.py
@@ -406,7 +406,7 @@ class TestAwsBaseHook:
"""
client_meta = AwsBaseHook(aws_conn_id=None, client_type="s3").conn_client_meta
- expected_user_agent_tag_keys = ["Airflow", "AmPP", "Caller", "DagRunKey"]
+ expected_user_agent_tag_keys = ["Airflow", "AmPP", "Caller", "DagRunKey", "MultiTeam", "TeamNameKey"]
result_user_agent_tags = client_meta.config.user_agent.split(" ")
result_user_agent_tag_keys = [tag.split("/")[0].lower() for tag in result_user_agent_tags]
@@ -473,6 +473,29 @@ class TestAwsBaseHook:
assert UUID(dag_run_key).version == expected_version
+ @pytest.mark.parametrize(
+ ("env_var", "expected"),
+ [
+ pytest.param({"AIRFLOW_CTX_TEAM_NAME": "my-team"}, "True", id="multi-team-enabled"),
+ pytest.param({}, "False", id="multi-team-disabled"),
+ ],
+ )
+ def test_user_agent_multi_team(self, env_var, expected):
+ with mock.patch.dict(os.environ, env_var, clear=True):
+ assert AwsBaseHook._is_multi_team() == (expected == "True")
+
+ @pytest.mark.parametrize(
+ ("env_var", "expected_version"),
+ [
+ pytest.param({"AIRFLOW_CTX_TEAM_NAME": "my-team"}, 5, id="team-set"),
+ pytest.param({}, None, id="team-not-set"),
+ ],
+ )
+ def test_user_agent_team_name_key_is_hashed_correctly(self, env_var, expected_version):
+ with mock.patch.dict(os.environ, env_var, clear=True):
+ team_name_key = AwsBaseHook._generate_team_name_key()
+ assert UUID(team_name_key).version == expected_version
+
@pytest.mark.parametrize(
"sts_endpoint",
[