This is an automated email from the ASF dual-hosted git repository. jedcunningham pushed a commit to branch v2-2-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit bda0a741f008344429f6e33bd7668a3b6f2ba6c8 Author: Ephraim Anierobi <[email protected]> AuthorDate: Mon Jan 17 13:45:51 2022 +0100 Add downgrade to some FAB migrations (#20874) There are some FAB migrations that don't have downgrades. This PR fixes it (cherry picked from commit f964fd47de5435f5c354c6211d24bd2dfaaffa9e) --- .../2c6edca13270_resource_based_permissions.py | 29 +++++++++ .../849da589634d_prefix_dag_permissions.py | 70 +++++++++++++++++++++- ...ad25_resource_based_permissions_for_default_.py | 29 +++++++++ 3 files changed, 127 insertions(+), 1 deletion(-) diff --git a/airflow/migrations/versions/2c6edca13270_resource_based_permissions.py b/airflow/migrations/versions/2c6edca13270_resource_based_permissions.py index 85673ce..02d7541 100644 --- a/airflow/migrations/versions/2c6edca13270_resource_based_permissions.py +++ b/airflow/migrations/versions/2c6edca13270_resource_based_permissions.py @@ -307,6 +307,31 @@ def remap_permissions(): appbuilder.sm.delete_action(old_action_name) +def undo_remap_permissions(): + """Unapply Map Airflow permissions""" + appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + for old, new in mapping.items: + (new_resource_name, new_action_name) = new[0] + new_permission = appbuilder.sm.get_permission(new_action_name, new_resource_name) + if not new_permission: + continue + for old_resource_name, old_action_name in old: + old_permission = appbuilder.sm.create_permission(old_action_name, old_resource_name) + for role in appbuilder.sm.get_all_roles(): + if appbuilder.sm.permission_exists_in_one_or_more_roles( + new_resource_name, new_action_name, [role.id] + ): + appbuilder.sm.add_permission_to_role(role, old_permission) + appbuilder.sm.remove_permission_from_role(role, new_permission) + appbuilder.sm.delete_permission(new_action_name, new_resource_name) + + if not appbuilder.sm.get_action(new_action_name): + continue + resources = appbuilder.sm.get_all_resources() + if not any(appbuilder.sm.get_permission(new_action_name, resource.name) for resource in resources): + appbuilder.sm.delete_action(new_action_name) + + def upgrade(): """Apply Resource based permissions.""" log = logging.getLogger() @@ -317,3 +342,7 @@ def upgrade(): def downgrade(): """Unapply Resource based permissions.""" + log = logging.getLogger() + handlers = log.handlers[:] + undo_remap_permissions() + log.handlers = handlers diff --git a/airflow/migrations/versions/849da589634d_prefix_dag_permissions.py b/airflow/migrations/versions/849da589634d_prefix_dag_permissions.py index fbbaa3a..adf25b1 100644 --- a/airflow/migrations/versions/849da589634d_prefix_dag_permissions.py +++ b/airflow/migrations/versions/849da589634d_prefix_dag_permissions.py @@ -55,6 +55,22 @@ def prefix_individual_dag_permissions(session): session.commit() +def remove_prefix_in_individual_dag_permissions(session): + dag_perms = ['can_read', 'can_edit'] + prefix = "DAG:" + perms = ( + session.query(Permission) + .join(Action) + .filter(Action.name.in_(dag_perms)) + .join(Resource) + .filter(Resource.name.like(prefix + '%')) + .all() + ) + for permission in perms: + permission.resource.name = permission.resource.name[len(prefix) :] + session.commit() + + def get_or_create_dag_resource(session): dag_resource = get_resource_query(session, permissions.RESOURCE_DAG).first() if dag_resource: @@ -68,6 +84,19 @@ def get_or_create_dag_resource(session): return dag_resource +def get_or_create_all_dag_resource(session): + all_dag_resource = get_resource_query(session, 'all_dags').first() + if all_dag_resource: + return all_dag_resource + + all_dag_resource = Resource() + all_dag_resource.name = 'all_dags' + session.add(all_dag_resource) + session.commit() + + return all_dag_resource + + def get_or_create_action(session, action_name): action = get_action_query(session, action_name).first() if action: @@ -158,6 +187,43 @@ def migrate_to_new_dag_permissions(db): db.session.commit() +def undo_migrate_to_new_dag_permissions(session): + # Remove prefix from individual dag perms + remove_prefix_in_individual_dag_permissions(session) + + # Update existing permissions to use `can_dag_read` instead of `can_read` + can_read_action = get_action_query(session, 'can_read').first() + new_can_read_permissions = get_permission_with_action_query(session, can_read_action) + can_dag_read_action = get_or_create_action(session, 'can_dag_read') + update_permission_action(session, new_can_read_permissions, can_dag_read_action) + + # Update existing permissions to use `can_dag_edit` instead of `can_edit` + can_edit_action = get_action_query(session, 'can_edit').first() + new_can_edit_permissions = get_permission_with_action_query(session, can_edit_action) + can_dag_edit_action = get_or_create_action(session, 'can_dag_edit') + update_permission_action(session, new_can_edit_permissions, can_dag_edit_action) + + # Update existing permissions for `DAGs` resource to use `all_dags` resource. + dag_resource = get_resource_query(session, permissions.RESOURCE_DAG).first() + if dag_resource: + new_dag_permission = get_permission_with_resource_query(session, dag_resource) + old_all_dag_resource = get_or_create_all_dag_resource(session) + update_permission_resource(session, new_dag_permission, old_all_dag_resource) + + # Delete the `DAG` resource + session.delete(dag_resource) + + # Delete `can_read` action + if can_read_action: + session.delete(can_read_action) + + # Delete `can_edit` action + if can_edit_action: + session.delete(can_edit_action) + + session.commit() + + def upgrade(): db = SQLA() db.session = settings.Session @@ -167,4 +233,6 @@ def upgrade(): def downgrade(): - pass + db = SQLA() + db.session = settings.Session + undo_migrate_to_new_dag_permissions(db.session) diff --git a/airflow/migrations/versions/a13f7613ad25_resource_based_permissions_for_default_.py b/airflow/migrations/versions/a13f7613ad25_resource_based_permissions_for_default_.py index 917bb8d..696c8ba 100644 --- a/airflow/migrations/versions/a13f7613ad25_resource_based_permissions_for_default_.py +++ b/airflow/migrations/versions/a13f7613ad25_resource_based_permissions_for_default_.py @@ -159,6 +159,31 @@ def remap_permissions(): appbuilder.sm.delete_action(old_action_name) +def undo_remap_permissions(): + """Unapply Map Airflow permissions""" + appbuilder = create_app(config={'FAB_UPDATE_PERMS': False}).appbuilder + for old, new in mapping.items(): + (new_resource_name, new_action_name) = new[0] + new_permission = appbuilder.sm.get_permission(new_action_name, new_resource_name) + if not new_permission: + continue + for old_action_name, old_resource_name in old: + old_permission = appbuilder.sm.create_permission(old_action_name, old_resource_name) + for role in appbuilder.sm.get_all_roles(): + if appbuilder.sm.permission_exists_in_one_or_more_roles( + new_resource_name, new_action_name, [role.id] + ): + appbuilder.sm.add_permission_to_role(role, old_permission) + appbuilder.sm.remove_permission_from_role(role, new_permission) + appbuilder.sm.delete_permission(new_action_name, new_resource_name) + + if not appbuilder.sm.get_action(new_action_name): + continue + resources = appbuilder.sm.get_all_resources() + if not any(appbuilder.sm.get_permission(new_action_name, resource.name) for resource in resources): + appbuilder.sm.delete_action(new_action_name) + + def upgrade(): """Apply Resource based permissions.""" log = logging.getLogger() @@ -169,3 +194,7 @@ def upgrade(): def downgrade(): """Unapply Resource based permissions.""" + log = logging.getLogger() + handlers = log.handlers[:] + undo_remap_permissions() + log.handlers = handlers
