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 dad2b927f1b Fix AWS auth manager batch team context (#66203)
dad2b927f1b is described below
commit dad2b927f1b73c7c42ba4177dd15ac7befa16476
Author: suyua9 <[email protected]>
AuthorDate: Tue May 5 00:20:26 2026 +0800
Fix AWS auth manager batch team context (#66203)
---
.../amazon/aws/auth_manager/aws_auth_manager.py | 12 ++++++
.../aws/auth_manager/test_aws_auth_manager.py | 45 +++++++++++++++++++---
2 files changed, 51 insertions(+), 6 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py
index 57febc258d0..03bd4c79b0b 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py
@@ -288,6 +288,9 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"entity_id": cast("ConnectionDetails",
request["details"]).conn_id
if request.get("details")
else None,
+ "team_name": self._get_team_name(cast("ConnectionDetails",
request["details"]))
+ if request.get("details")
+ else None,
},
)
for request in requests
@@ -309,6 +312,9 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"entity_id": cast("DagDetails", request["details"]).id
if request.get("details")
else None,
+ "team_name": self._get_team_name(cast("DagDetails",
request["details"]))
+ if request.get("details")
+ else None,
"context": {
"dag_entity": {
"string": cast("DagAccessEntity",
request["access_entity"]).value,
@@ -337,6 +343,9 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"entity_id": cast("PoolDetails", request["details"]).name
if request.get("details")
else None,
+ "team_name": self._get_team_name(cast("PoolDetails",
request["details"]))
+ if request.get("details")
+ else None,
},
)
for request in requests
@@ -358,6 +367,9 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"entity_id": cast("VariableDetails",
request["details"]).key
if request.get("details")
else None,
+ "team_name": self._get_team_name(cast("VariableDetails",
request["details"]))
+ if request.get("details")
+ else None,
},
)
for request in requests
diff --git
a/providers/amazon/tests/unit/amazon/aws/auth_manager/test_aws_auth_manager.py
b/providers/amazon/tests/unit/amazon/aws/auth_manager/test_aws_auth_manager.py
index 6d6bb72ac32..91d69c5bf19 100644
---
a/providers/amazon/tests/unit/amazon/aws/auth_manager/test_aws_auth_manager.py
+++
b/providers/amazon/tests/unit/amazon/aws/auth_manager/test_aws_auth_manager.py
@@ -64,6 +64,12 @@ else:
mock = Mock()
+
+def _with_team_name(details, team_name="team_a"):
+ setattr(details, "team_name", team_name)
+ return details
+
+
SAML_METADATA_PARSED = {
"idp": {
"entityId":
"https://portal.sso.us-east-1.amazonaws.com/saml/assertion/<assertion>",
@@ -447,7 +453,9 @@ class TestAwsAuthManager:
mock_avp_facade.is_authorized = is_authorized
result = auth_manager.is_authorized_dag(
- method="GET", details=DagDetails(id="dag_1", team_name="team_a"),
user=mock
+ method="GET",
+ details=_with_team_name(DagDetails(id="dag_1")),
+ user=mock,
)
is_authorized.assert_called_once_with(
@@ -588,7 +596,10 @@ class TestAwsAuthManager:
result = auth_manager.batch_is_authorized_connection(
requests=[
{"method": "GET"},
- {"method": "PUT", "details":
ConnectionDetails(conn_id="test")},
+ {
+ "method": "PUT",
+ "details":
_with_team_name(ConnectionDetails(conn_id="test")),
+ },
],
user=mock,
)
@@ -599,11 +610,13 @@ class TestAwsAuthManager:
"method": "GET",
"entity_type": AvpEntities.CONNECTION,
"entity_id": None,
+ "team_name": None,
},
{
"method": "PUT",
"entity_type": AvpEntities.CONNECTION,
"entity_id": "test",
+ "team_name": "team_a",
},
],
user=ANY,
@@ -622,10 +635,17 @@ class TestAwsAuthManager:
result = auth_manager.batch_is_authorized_dag(
requests=[
{"method": "GET"},
- {"method": "GET", "details": DagDetails(id="dag_1")},
+ {
+ "method": "GET",
+ "details": _with_team_name(DagDetails(id="dag_1")),
+ },
]
+ [
- {"method": "GET", "details": DagDetails(id="dag_1"),
"access_entity": dag_access_entity}
+ {
+ "method": "GET",
+ "details": _with_team_name(DagDetails(id="dag_1")),
+ "access_entity": dag_access_entity,
+ }
for dag_access_entity in (
DagAccessEntity.AUDIT_LOG,
DagAccessEntity.CODE,
@@ -648,12 +668,14 @@ class TestAwsAuthManager:
"method": "GET",
"entity_type": AvpEntities.DAG,
"entity_id": None,
+ "team_name": None,
"context": None,
},
{
"method": "GET",
"entity_type": AvpEntities.DAG,
"entity_id": "dag_1",
+ "team_name": "team_a",
"context": None,
},
]
@@ -662,6 +684,7 @@ class TestAwsAuthManager:
"method": "GET",
"entity_type": AvpEntities.DAG,
"entity_id": "dag_1",
+ "team_name": "team_a",
"context": {"dag_entity": {"string": dag_entity}},
}
for dag_entity in (
@@ -693,7 +716,10 @@ class TestAwsAuthManager:
result = auth_manager.batch_is_authorized_pool(
requests=[
{"method": "GET"},
- {"method": "PUT", "details": PoolDetails(name="test")},
+ {
+ "method": "PUT",
+ "details": _with_team_name(PoolDetails(name="test")),
+ },
],
user=mock,
)
@@ -704,11 +730,13 @@ class TestAwsAuthManager:
"method": "GET",
"entity_type": AvpEntities.POOL,
"entity_id": None,
+ "team_name": None,
},
{
"method": "PUT",
"entity_type": AvpEntities.POOL,
"entity_id": "test",
+ "team_name": "team_a",
},
],
user=ANY,
@@ -727,7 +755,10 @@ class TestAwsAuthManager:
result = auth_manager.batch_is_authorized_variable(
requests=[
{"method": "GET"},
- {"method": "PUT", "details": VariableDetails(key="test")},
+ {
+ "method": "PUT",
+ "details": _with_team_name(VariableDetails(key="test")),
+ },
],
user=mock,
)
@@ -738,11 +769,13 @@ class TestAwsAuthManager:
"method": "GET",
"entity_type": AvpEntities.VARIABLE,
"entity_id": None,
+ "team_name": None,
},
{
"method": "PUT",
"entity_type": AvpEntities.VARIABLE,
"entity_id": "test",
+ "team_name": "team_a",
},
],
user=ANY,
