Skip to content

Commit 691c143

Browse files
vincbeckvincbeck
committed
Change menu action in AWS auth manager
1 parent d4aa0e5 commit 691c143

4 files changed

Lines changed: 44 additions & 48 deletions

File tree

‎airflow/providers/amazon/aws/auth_manager/avp/entities.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ class AvpEntities(Enum):
3838
CUSTOM = "Custom"
3939
DAG = "Dag"
4040
DATASET = "Dataset"
41+
MENU = "Menu"
4142
POOL = "Pool"
4243
VARIABLE = "Variable"
4344
VIEW = "View"

‎airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py‎

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -513,17 +513,13 @@ def get_cli_commands() -> list[CLICommand]:
513513
),
514514
]
515515

516-
def _get_menu_item_request(self, resource_name: str) -> IsAuthorizedRequest:
517-
menu_item_request = _MENU_ITEM_REQUESTS.get(resource_name)
518-
if menu_item_request:
519-
return menu_item_request
520-
else:
521-
self.log.info("The menu item '%s' is unknown. It must come from a plugin", resource_name)
522-
return {
523-
"method": "MENU",
524-
"entity_type": AvpEntities.CUSTOM,
525-
"entity_id": resource_name,
526-
}
516+
@staticmethod
517+
def _get_menu_item_request(resource_name: str) -> IsAuthorizedRequest:
518+
return {
519+
"method": "MENU",
520+
"entity_type": AvpEntities.MENU,
521+
"entity_id": resource_name,
522+
}
527523

528524
def _has_access_to_menu_item(
529525
self, batch_is_authorized_results: list[dict], request: IsAuthorizedRequest, user: AwsAuthManagerUser

‎airflow/providers/amazon/aws/auth_manager/cli/schema.json‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,12 @@
121121
"resourceTypes": ["Dataset"]
122122
}
123123
},
124+
"Menu.MENU": {
125+
"appliesTo": {
126+
"principalTypes": ["User"],
127+
"resourceTypes": ["Menu"]
128+
}
129+
},
124130
"Pool.DELETE": {
125131
"appliesTo": {
126132
"principalTypes": ["User"],
@@ -182,6 +188,7 @@
182188
"Custom": {},
183189
"Dag": {},
184190
"Dataset": {},
191+
"Menu": {},
185192
"Pool": {},
186193
"Role": {},
187194
"User": {

‎tests/providers/amazon/aws/auth_manager/test_aws_auth_manager.py‎

Lines changed: 29 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -523,55 +523,48 @@ def test_filter_permitted_menu_items(self, mock_get_user, auth_manager, test_use
523523
{
524524
"request": {
525525
"principal": {"entityType": "Airflow::User", "entityId": "test_user_id"},
526-
"action": {"actionType": "Airflow::Action", "actionId": "Connection.MENU"},
527-
"resource": {"entityType": "Airflow::Connection", "entityId": "*"},
526+
"action": {"actionType": "Airflow::Action", "actionId": "Menu.MENU"},
527+
"resource": {"entityType": "Airflow::Menu", "entityId": "Connections"},
528528
},
529529
"decision": "DENY",
530530
},
531531
{
532532
"request": {
533533
"principal": {"entityType": "Airflow::User", "entityId": "test_user_id"},
534-
"action": {"actionType": "Airflow::Action", "actionId": "Variable.MENU"},
535-
"resource": {"entityType": "Airflow::Variable", "entityId": "*"},
534+
"action": {"actionType": "Airflow::Action", "actionId": "Menu.MENU"},
535+
"resource": {"entityType": "Airflow::Menu", "entityId": "Variables"},
536536
},
537537
"decision": "ALLOW",
538538
},
539539
{
540540
"request": {
541541
"principal": {"entityType": "Airflow::User", "entityId": "test_user_id"},
542-
"action": {"actionType": "Airflow::Action", "actionId": "Dataset.MENU"},
543-
"resource": {"entityType": "Airflow::Dataset", "entityId": "*"},
542+
"action": {"actionType": "Airflow::Action", "actionId": "Menu.MENU"},
543+
"resource": {"entityType": "Airflow::Menu", "entityId": "Datasets"},
544544
},
545545
"decision": "DENY",
546546
},
547547
{
548548
"request": {
549549
"principal": {"entityType": "Airflow::User", "entityId": "test_user_id"},
550-
"action": {"actionType": "Airflow::Action", "actionId": "View.MENU"},
551-
"resource": {"entityType": "Airflow::View", "entityId": "CLUSTER_ACTIVITY"},
550+
"action": {"actionType": "Airflow::Action", "actionId": "Menu.MENU"},
551+
"resource": {"entityType": "Airflow::Menu", "entityId": "Cluster Activity"},
552552
},
553553
"decision": "DENY",
554554
},
555555
{
556556
"request": {
557557
"principal": {"entityType": "Airflow::User", "entityId": "test_user_id"},
558-
"action": {"actionType": "Airflow::Action", "actionId": "Dag.MENU"},
559-
"resource": {"entityType": "Airflow::Dag", "entityId": "*"},
560-
"context": {
561-
"contextMap": {
562-
"dag_entity": {
563-
"string": "AUDIT_LOG",
564-
}
565-
}
566-
},
558+
"action": {"actionType": "Airflow::Action", "actionId": "Menu.MENU"},
559+
"resource": {"entityType": "Airflow::Menu", "entityId": "Audit Logs"},
567560
},
568561
"decision": "ALLOW",
569562
},
570563
{
571564
"request": {
572565
"principal": {"entityType": "Airflow::User", "entityId": "test_user_id"},
573-
"action": {"actionType": "Airflow::Action", "actionId": "Custom.MENU"},
574-
"resource": {"entityType": "Airflow::Custom", "entityId": "CustomPage"},
566+
"action": {"actionType": "Airflow::Action", "actionId": "Menu.MENU"},
567+
"resource": {"entityType": "Airflow::Menu", "entityId": "CustomPage"},
575568
},
576569
"decision": "ALLOW",
577570
},
@@ -592,37 +585,36 @@ def test_filter_permitted_menu_items(self, mock_get_user, auth_manager, test_use
592585
]
593586
)
594587

588+
"""
589+
return {
590+
"method": "MENU",
591+
"entity_type": AvpEntities.MENU,
592+
"entity_id": resource_name,
593+
}
594+
"""
595+
595596
auth_manager.avp_facade.get_batch_is_authorized_results.assert_called_once_with(
596597
requests=[
597598
{
598599
"method": "MENU",
599-
"entity_type": AvpEntities.CONNECTION,
600-
},
601-
{
602-
"method": "MENU",
603-
"entity_type": AvpEntities.VARIABLE,
600+
"entity_type": AvpEntities.MENU,
601+
"entity_id": "Connections",
604602
},
605603
{
606604
"method": "MENU",
607-
"entity_type": AvpEntities.DATASET,
605+
"entity_type": AvpEntities.MENU,
606+
"entity_id": "Variables",
608607
},
609608
{
610609
"method": "MENU",
611-
"entity_type": AvpEntities.VIEW,
612-
"entity_id": AccessView.CLUSTER_ACTIVITY.value,
613-
},
614-
{
615-
"method": "MENU",
616-
"entity_type": AvpEntities.DAG,
617-
"context": {
618-
"dag_entity": {
619-
"string": DagAccessEntity.AUDIT_LOG.value,
620-
},
621-
},
610+
"entity_type": AvpEntities.MENU,
611+
"entity_id": "Datasets",
622612
},
613+
{"method": "MENU", "entity_type": AvpEntities.MENU, "entity_id": "Cluster Activity"},
614+
{"method": "MENU", "entity_type": AvpEntities.MENU, "entity_id": "Audit Logs"},
623615
{
624616
"method": "MENU",
625-
"entity_type": AvpEntities.CUSTOM,
617+
"entity_type": AvpEntities.MENU,
626618
"entity_id": "CustomPage",
627619
},
628620
],

0 commit comments

Comments
 (0)
⚡