├── README.md └── rbac_roles_cli.py /README.md: -------------------------------------------------------------------------------- 1 | # airflow-rbac-roles-cli 2 | Airflow provides a way to manage permissions: RBAC. However, this system is only available from UI. Many use cases such as the implementation of an Airflow orchestrator for multiple projects need to take profit of dag-level permissions in order to only authorize some users to access specific project dags. **Creating roles with dag-level permissions on multiple dags directly from the UI is very error-prone and time-consuming.** 3 | 4 | That's why I created a script which given the **airflow url**, the **dag names** and the **new role name**, will automatically create the role with the needed permissions to access only these dags and not the other dags deployed on Airflow. 5 | 6 | > The tool can be used to create roles for multiple projects inside a same Airflow. 7 | 8 | ## Usage 9 | ```shell 10 | python3 rbac_roles_cli.py -u airflow_url -r new_role_name -d dag1 dag2 11 | ``` 12 | Will create a role named `new_role_name` on the airflow project at the url `airflow_url` with enough permissions to only manage the dags `dag1` and `dag2`. 13 | 14 | Note that the permissions can be edited directly in the code and following the syntax used in the Airflow code. 15 | 16 | ## Compatibility with Cloud Composer 17 | GCP offers Cloud Composer: a Google-managed Airflow. The script also work with Composer, you only need to add an argument to the command: the Google access-token. This token is used in order to access the Composer page which is protected by a Google authentication page. 18 | 19 | ```shell 20 | python3 rbac_roles_cli.py -u airflow_url -r new_role_name -d dag1 dag2 -t access_token 21 | ``` 22 | -------------------------------------------------------------------------------- /rbac_roles_cli.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | import logging 3 | import requests 4 | import argparse 5 | import base64 6 | 7 | def create_rbac_role_with_permissions( 8 | airflow_url: str, 9 | new_role_name: str, 10 | dag_names: List[str], 11 | google_access_token: str=None, 12 | airflow_username: str=None, 13 | airflow_password: str=None 14 | ): 15 | headers = { 16 | "Accept": "application/json", 17 | "Content-Type": "application/json", 18 | } 19 | 20 | if google_access_token: 21 | headers["Authorization"] = "Bearer " + google_access_token 22 | elif airflow_username and airflow_password: 23 | auth_str = f"{airflow_username}:{airflow_password}" 24 | base64_auth_str = base64.b64encode(auth_str.encode()).decode() 25 | headers["Authorization"] = "Basic " + base64_auth_str 26 | 27 | read = "can_read" 28 | edit = "can_edit" 29 | create = "can_create" 30 | delete = "can_delete" 31 | menu = "menu_access" 32 | 33 | # add general permissions 34 | permissions = [] 35 | read_permissions = make_permissions(read,["Task Instances", "Website", "DAG Runs", "Audit Logs", "ImportError", "XComs", \ 36 | "DAG Code", "Plugins", "My Password", "My Profile", "Jobs", "SLA Misses", "DAG Dependencies", "Task Logs"]) 37 | edit_permissions = make_permissions(edit, ["Task Instances", "My Password", "My Profile", "DAG Runs"]) 38 | create_permissions = make_permissions(create, ["DAG Runs", "Task Instances"]) 39 | delete_permissions = make_permissions(delete, ["DAG Runs", "Task Instances"]) 40 | menu_permissions = make_permissions(menu, ["View Menus", "Browse", "Docs", "Documentation", "SLA Misses", "Jobs", "DAG Runs", \ 41 | "Audit Logs", "Task Instances", "DAG Dependencies"]) 42 | permissions += read_permissions + edit_permissions + create_permissions + delete_permissions + menu_permissions 43 | 44 | # add dag-specific permissions 45 | for dag in dag_names: 46 | dag = "DAG:" + dag 47 | read_permissions = make_permissions(read,[dag]) 48 | edit_permissions = make_permissions(edit, [dag]) 49 | delete_permissions = make_permissions(delete, [dag]) 50 | permissions += read_permissions + edit_permissions + delete_permissions 51 | 52 | data = { 53 | "actions": [ 54 | *permissions 55 | ], 56 | "name": new_role_name 57 | } 58 | 59 | airflow_url += "/api/v1/roles" 60 | response = requests.post(airflow_url, json=data, headers=headers) 61 | 62 | if response.status_code == 403: 63 | raise PermissionError(f"Error 403 returned, please check if your AirFlow account is Op/Admin or verify the dags exist. \n {response.json()}") 64 | elif response.status_code == 401: 65 | raise PermissionError(f"Error 401 returned, please check the access token if the page is protected by an authentication") 66 | elif response.status_code == 200: 67 | print(f"Role `{new_role_name}` successfully created.") 68 | return 69 | elif response.status_code == 409: # Role already exists, update it 70 | print("Role already exists, updating...") 71 | airflow_role_update_url = f"{airflow_url}/{new_role_name}" 72 | update_response = requests.patch(airflow_role_update_url, json=data, headers=headers) 73 | if update_response.status_code == 200: 74 | print(f"Role `{new_role_name}` successfully updated.") 75 | else: 76 | raise ConnectionError(f"An error occurred during role update: {update_response.json()}") 77 | else: 78 | raise ConnectionError(f"An error occurred during role creation: {response.json()}") 79 | 80 | def make_permissions(action, resources): 81 | permissions = [] 82 | for perm in resources: 83 | permissions.append(make_permission(action, perm)) 84 | return permissions 85 | 86 | def make_permission(action, resource): 87 | return { 88 | "action": {"name": action}, 89 | "resource": {"name": resource} 90 | } 91 | 92 | 93 | if __name__ == "__main__": 94 | parser = argparse.ArgumentParser() 95 | parser.add_argument("-u", "--airflow-url", required=True, help="URL to the composer Airflow UI root page") 96 | parser.add_argument("-r", "--role-name", required=True, help="Name of the new created role") 97 | parser.add_argument("-d", "--dags", nargs="+", required=True, help="List of accessible dags for the role") 98 | parser.add_argument("-t", "--access-token", required=False, help="Google access token used only if Airflow is managed by Cloud Composer") 99 | parser.add_argument("-afu", "--airflow-username", required=False, help="Airflow username for Basic Auth") 100 | parser.add_argument("-afp", "--airflow-password", required=False, help="Airflow password for Basic Auth") 101 | 102 | args = parser.parse_args() 103 | create_rbac_role_with_permissions( 104 | airflow_url=args.airflow_url, 105 | new_role_name=args.role_name, 106 | dag_names=args.dags, 107 | google_access_token=args.access_token, 108 | airflow_username=args.airflow_username, 109 | airflow_password=args.airflow_password 110 | ) 111 | --------------------------------------------------------------------------------