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 551870b59dd Update `AwsAuthManager` to support multi-team #65371
(#65393)
551870b59dd is described below
commit 551870b59dd671d73c3b15cb02aa40be9c55be65
Author: Haseeb Malik <[email protected]>
AuthorDate: Tue Apr 21 10:35:18 2026 -0400
Update `AwsAuthManager` to support multi-team #65371 (#65393)
---
.../amazon/aws/auth_manager/avp/entities.py | 1 +
.../amazon/aws/auth_manager/avp/facade.py | 22 ++-
.../amazon/aws/auth_manager/avp/schema.json | 198 +++++++++++++++++++--
.../amazon/aws/auth_manager/aws_auth_manager.py | 30 ++++
.../aws/auth_manager/test_aws_auth_manager.py | 132 +++++++++++++-
5 files changed, 364 insertions(+), 19 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/entities.py
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/entities.py
index 87e9ed70bd1..6894fad6255 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/entities.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/entities.py
@@ -42,6 +42,7 @@ class AvpEntities(Enum):
DAG = "Dag"
MENU = "Menu"
POOL = "Pool"
+ TEAM = "Team"
VARIABLE = "Variable"
VIEW = "View"
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/facade.py
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/facade.py
index a6de7be60b7..ac836690ce0 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/facade.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/facade.py
@@ -58,6 +58,7 @@ class IsAuthorizedRequest(TypedDict, total=False):
entity_type: AvpEntities
entity_id: str | None
context: dict | None
+ team_name: str | None
class AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
@@ -87,6 +88,7 @@ class
AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
user: AwsAuthManagerUser,
entity_id: str | None = None,
context: dict | None = None,
+ team_name: str | None = None,
) -> bool:
"""
Make an authorization decision against Amazon Verified Permissions.
@@ -102,7 +104,9 @@ class
AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
:param entity_id: the entity ID the user accesses. If not provided,
all entities of the type will be
considered.
:param context: optional additional context to pass to Amazon Verified
Permissions.
+ :param team_name: optional team name for multi-team authorization.
"""
+ context = self._enrich_context_with_team(context, team_name)
entity_list = self._get_user_group_entities(user)
self.log.debug(
@@ -259,6 +263,21 @@ class
AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
]
return [user_entity, *group_entities]
+ @staticmethod
+ def _enrich_context_with_team(context: dict | None, team_name: str | None)
-> dict | None:
+ """
+ Add team_name to context if provided.
+
+ :param context: existing context dict or None
+ :param team_name: team name to add to context
+ """
+ if not team_name:
+ return context
+ team_context = {"team_name": {"string": team_name}}
+ if context is None:
+ return team_context
+ return {**context, **team_context}
+
@staticmethod
def _build_context(context: dict | None) -> dict | None:
if context is None or len(context) == 0:
@@ -275,6 +294,7 @@ class
AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
:param request: the request information
:param user: the user
"""
+ context = self._enrich_context_with_team(request.get("context"),
request.get("team_name"))
return prune_dict(
{
"principal": {"entityType": get_entity_type(AvpEntities.USER),
"entityId": user.get_id()},
@@ -288,6 +308,6 @@ class
AwsAuthManagerAmazonVerifiedPermissionsFacade(LoggingMixin):
"entityType": get_entity_type(request["entity_type"]),
"entityId": request.get("entity_id", "*"),
},
- "context": self._build_context(request.get("context")),
+ "context": self._build_context(context),
}
)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/schema.json
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/schema.json
index a90b9197115..cef4912bb4f 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/schema.json
+++
b/providers/amazon/src/airflow/providers/amazon/aws/auth_manager/avp/schema.json
@@ -64,31 +64,76 @@
"Connection.DELETE": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Connection"]
+ "resourceTypes": ["Connection"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Connection.GET": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Connection"]
+ "resourceTypes": ["Connection"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Connection.LIST": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Connection"]
+ "resourceTypes": ["Connection"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Connection.POST": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Connection"]
+ "resourceTypes": ["Connection"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Connection.PUT": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Connection"]
+ "resourceTypes": ["Connection"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Custom.DELETE": {
@@ -142,6 +187,10 @@
"dag_entity": {
"type": "String",
"required": false
+ },
+ "team_name": {
+ "type": "String",
+ "required": false
}
},
"type": "Record"
@@ -157,6 +206,10 @@
"dag_entity": {
"required": false,
"type": "String"
+ },
+ "team_name": {
+ "required": false,
+ "type": "String"
}
},
"type": "Record"
@@ -172,6 +225,10 @@
"dag_entity": {
"required": false,
"type": "String"
+ },
+ "team_name": {
+ "required": false,
+ "type": "String"
}
},
"type": "Record"
@@ -187,6 +244,10 @@
"dag_entity": {
"required": false,
"type": "String"
+ },
+ "team_name": {
+ "required": false,
+ "type": "String"
}
},
"type": "Record"
@@ -202,6 +263,10 @@
"dag_entity": {
"required": false,
"type": "String"
+ },
+ "team_name": {
+ "required": false,
+ "type": "String"
}
},
"type": "Record"
@@ -217,61 +282,163 @@
"Pool.DELETE": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Pool"]
+ "resourceTypes": ["Pool"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Pool.GET": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Pool"]
+ "resourceTypes": ["Pool"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Pool.LIST": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Pool"]
+ "resourceTypes": ["Pool"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Pool.POST": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Pool"]
+ "resourceTypes": ["Pool"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Pool.PUT": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Pool"]
+ "resourceTypes": ["Pool"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
+ }
+ },
+ "Team.GET": {
+ "appliesTo": {
+ "principalTypes": ["User"],
+ "resourceTypes": ["Team"]
+ }
+ },
+ "Team.LIST": {
+ "appliesTo": {
+ "principalTypes": ["User"],
+ "resourceTypes": ["Team"]
}
},
"Variable.DELETE": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Variable"]
+ "resourceTypes": ["Variable"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Variable.GET": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Variable"]
+ "resourceTypes": ["Variable"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Variable.LIST": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Variable"]
+ "resourceTypes": ["Variable"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Variable.POST": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Variable"]
+ "resourceTypes": ["Variable"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"Variable.PUT": {
"appliesTo": {
"principalTypes": ["User"],
- "resourceTypes": ["Variable"]
+ "resourceTypes": ["Variable"],
+ "context": {
+ "attributes": {
+ "team_name": {
+ "type": "String",
+ "required": false
+ }
+ },
+ "type": "Record"
+ }
}
},
"View.GET": {
@@ -291,6 +458,7 @@
"Dag": {},
"Menu": {},
"Pool": {},
+ "Team": {},
"Group": {},
"User": {
"memberOfTypes": [
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 f20f4f0a075..57febc258d0 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
@@ -56,6 +56,7 @@ if TYPE_CHECKING:
DagAccessEntity,
DagDetails,
PoolDetails,
+ TeamDetails,
VariableDetails,
)
from airflow.api_fastapi.common.types import MenuItem
@@ -125,11 +126,13 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
details: ConnectionDetails | None = None,
) -> bool:
connection_id = details.conn_id if details else None
+ team_name = self._get_team_name(details)
return self.avp_facade.is_authorized(
method=method,
entity_type=AvpEntities.CONNECTION,
user=user,
entity_id=connection_id,
+ team_name=team_name,
)
def is_authorized_dag(
@@ -141,6 +144,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
details: DagDetails | None = None,
) -> bool:
dag_id = details.id if details else None
+ team_name = self._get_team_name(details)
context = (
None
if access_entity is None
@@ -156,6 +160,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
user=user,
entity_id=dag_id,
context=context,
+ team_name=team_name,
)
def is_authorized_backfill(
@@ -193,22 +198,37 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
self, *, method: ResourceMethod, user: AwsAuthManagerUser, details:
PoolDetails | None = None
) -> bool:
pool_name = details.name if details else None
+ team_name = self._get_team_name(details)
return self.avp_facade.is_authorized(
method=method,
entity_type=AvpEntities.POOL,
user=user,
entity_id=pool_name,
+ team_name=team_name,
)
def is_authorized_variable(
self, *, method: ResourceMethod, user: AwsAuthManagerUser, details:
VariableDetails | None = None
) -> bool:
variable_key = details.key if details else None
+ team_name = self._get_team_name(details)
return self.avp_facade.is_authorized(
method=method,
entity_type=AvpEntities.VARIABLE,
user=user,
entity_id=variable_key,
+ team_name=team_name,
+ )
+
+ def is_authorized_team(
+ self, *, method: ResourceMethod, user: AwsAuthManagerUser, details:
TeamDetails | None = None
+ ) -> bool:
+ team_name = details.name if details else None
+ return self.avp_facade.is_authorized(
+ method=method,
+ entity_type=AvpEntities.TEAM,
+ user=user,
+ entity_id=team_name,
)
def is_authorized_view(
@@ -359,6 +379,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"method": method,
"entity_type": AvpEntities.CONNECTION,
"entity_id": conn_id,
+ "team_name": team_name,
}
requests[conn_id][method] = request
requests_list.append(request)
@@ -390,6 +411,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"method": method,
"entity_type": AvpEntities.DAG,
"entity_id": dag_id,
+ "team_name": team_name,
}
requests[dag_id][method] = request
requests_list.append(request)
@@ -421,6 +443,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"method": method,
"entity_type": AvpEntities.POOL,
"entity_id": pool_name,
+ "team_name": team_name,
}
requests[pool_name][method] = request
requests_list.append(request)
@@ -452,6 +475,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
"method": method,
"entity_type": AvpEntities.VARIABLE,
"entity_id": variable_key,
+ "team_name": team_name,
}
requests[variable_key][method] = request
requests_list.append(request)
@@ -493,6 +517,12 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
return app
+ @staticmethod
+ def _get_team_name(
+ details: ConnectionDetails | DagDetails | PoolDetails |
VariableDetails | None,
+ ) -> str | None:
+ return getattr(details, "team_name", None) if details else None
+
@staticmethod
def _get_menu_item_request(menu_item_text: str) -> IsAuthorizedRequest:
return {
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 8e0ff06c185..6d6bb72ac32 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
@@ -24,7 +24,7 @@ import pytest
from airflow.exceptions import AirflowProviderDeprecationWarning
-from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS,
AIRFLOW_V_3_2_PLUS
if not AIRFLOW_V_3_0_PLUS:
pytest.skip("AWS auth manager is only compatible with Airflow >= 3.0.0",
allow_module_level=True)
@@ -40,6 +40,11 @@ from
airflow.api_fastapi.auth.managers.models.resource_details import (
PoolDetails,
VariableDetails,
)
+
+if AIRFLOW_V_3_2_PLUS:
+ from airflow.api_fastapi.auth.managers.models.resource_details import
TeamDetails
+else:
+ TeamDetails = None # type: ignore[assignment,misc]
from airflow.api_fastapi.common.types import MenuItem
from airflow.providers.amazon.aws.auth_manager.avp.entities import AvpEntities
from airflow.providers.amazon.aws.auth_manager.aws_auth_manager import
AwsAuthManager
@@ -158,6 +163,7 @@ class TestAwsAuthManager:
entity_type=AvpEntities.CONNECTION,
user=expected_user,
entity_id=expected_entity_id,
+ team_name=None,
)
assert result
@@ -206,6 +212,7 @@ class TestAwsAuthManager:
user=expected_user,
entity_id=expected_entity_id,
context=expected_context,
+ team_name=None,
)
assert result
@@ -328,7 +335,11 @@ class TestAwsAuthManager:
result = auth_manager.is_authorized_pool(method=method,
details=details, user=user)
is_authorized.assert_called_once_with(
- method=method, entity_type=AvpEntities.POOL, user=expected_user,
entity_id=expected_entity_id
+ method=method,
+ entity_type=AvpEntities.POOL,
+ user=expected_user,
+ entity_id=expected_entity_id,
+ team_name=None,
)
assert result
@@ -356,7 +367,122 @@ class TestAwsAuthManager:
result = auth_manager.is_authorized_variable(method=method,
details=details, user=user)
is_authorized.assert_called_once_with(
- method=method, entity_type=AvpEntities.VARIABLE,
user=expected_user, entity_id=expected_entity_id
+ method=method,
+ entity_type=AvpEntities.VARIABLE,
+ user=expected_user,
+ entity_id=expected_entity_id,
+ team_name=None,
+ )
+ assert result
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="TeamDetails not
available before Airflow 3.2.0")
+ @pytest.mark.parametrize(
+ (
+ "function",
+ "details_cls",
+ "details_kwargs",
+ "entity_type",
+ "expected_entity_id",
+ "expected_team_name",
+ ),
+ [
+ (
+ "is_authorized_connection",
+ ConnectionDetails,
+ {"conn_id": "conn_id", "team_name": "team_a"},
+ AvpEntities.CONNECTION,
+ "conn_id",
+ "team_a",
+ ),
+ (
+ "is_authorized_pool",
+ PoolDetails,
+ {"name": "pool1", "team_name": "team_a"},
+ AvpEntities.POOL,
+ "pool1",
+ "team_a",
+ ),
+ (
+ "is_authorized_variable",
+ VariableDetails,
+ {"key": "var1", "team_name": "team_a"},
+ AvpEntities.VARIABLE,
+ "var1",
+ "team_a",
+ ),
+ ],
+ )
+ @patch.object(AwsAuthManager, "avp_facade")
+ def test_is_authorized_with_team_name(
+ self,
+ mock_avp_facade,
+ function,
+ details_cls,
+ details_kwargs,
+ entity_type,
+ expected_entity_id,
+ expected_team_name,
+ auth_manager,
+ ):
+ is_authorized = Mock(return_value=True)
+ mock_avp_facade.is_authorized = is_authorized
+ details = details_cls(**details_kwargs)
+
+ method: ResourceMethod = "GET"
+ result = getattr(auth_manager, function)(method=method,
details=details, user=mock)
+
+ is_authorized.assert_called_once_with(
+ method=method,
+ entity_type=entity_type,
+ user=mock,
+ entity_id=expected_entity_id,
+ team_name=expected_team_name,
+ )
+ assert result
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="TeamDetails not
available before Airflow 3.2.0")
+ @patch.object(AwsAuthManager, "avp_facade")
+ def test_is_authorized_dag_with_team_name(self, mock_avp_facade,
auth_manager):
+ is_authorized = Mock(return_value=True)
+ 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
+ )
+
+ is_authorized.assert_called_once_with(
+ method="GET",
+ entity_type=AvpEntities.DAG,
+ user=mock,
+ entity_id="dag_1",
+ context=None,
+ team_name="team_a",
+ )
+ assert result
+
+ @patch.object(AwsAuthManager, "avp_facade")
+ def test_is_authorized_team(self, mock_avp_facade, auth_manager):
+ is_authorized = Mock(return_value=True)
+ mock_avp_facade.is_authorized = is_authorized
+
+ result = auth_manager.is_authorized_team(method="GET", details=None,
user=mock)
+
+ is_authorized.assert_called_once_with(
+ method="GET", entity_type=AvpEntities.TEAM, user=mock,
entity_id=None
+ )
+ assert result
+
+ @pytest.mark.skipif(not AIRFLOW_V_3_2_PLUS, reason="TeamDetails not
available before Airflow 3.2.0")
+ @patch.object(AwsAuthManager, "avp_facade")
+ def test_is_authorized_team_with_details(self, mock_avp_facade,
auth_manager):
+ is_authorized = Mock(return_value=True)
+ mock_avp_facade.is_authorized = is_authorized
+ details = TeamDetails(name="team_a")
+
+ result = auth_manager.is_authorized_team(method="GET",
details=details, user=mock)
+
+ is_authorized.assert_called_once_with(
+ method="GET", entity_type=AvpEntities.TEAM, user=mock,
entity_id="team_a"
)
assert result
