[GitHub] [airflow] ashb commented on a change in pull request #12720: Refactor prefix migration

[GitBox <gi...@apache.org>]

ashb commented on a change in pull request #12720:
URL: https://github.com/apache/airflow/pull/12720#discussion_r532783651



##########
File path: airflow/migrations/versions/849da589634d_prefix_dag_permissions.py
##########
@@ -34,80 +37,139 @@
 depends_on = None
 
 
+def prefix_individual_dag_permissions(session):  # noqa: D103
+    dag_perms = ['can_dag_read', 'can_dag_edit']
+    prefix = "DAG:"
+    permission_view_menus = (
+        session.query(PermissionView)
+        .join(Permission)
+        .filter(Permission.name.in_(dag_perms))
+        .join(ViewMenu)
+        .filter(ViewMenu.name != 'all_dags')
+        .filter(ViewMenu.name.notlike('DAG:%'))
+        .all()
+    )
+    view_menu_ids = {pvm.view_menu.id for pvm in permission_view_menus}
+    vm_query = session.query(ViewMenu).filter(ViewMenu.id.in_(view_menu_ids))
+    vm_query.update({ViewMenu.name: prefix + ViewMenu.name}, synchronize_session=False)
+    session.commit()
+
+
+def get_or_create_dag_resource(session):  # noqa: D103
+    dag_resource = get_resource_query(session, permissions.RESOURCE_DAG).first()
+    if dag_resource:
+        return dag_resource
+
+    dag_resource = ViewMenu()
+    dag_resource.name = permissions.RESOURCE_DAG
+    session.add(dag_resource)
+    session.commit()
+
+    return dag_resource
+
+
+def get_or_create_action(session, action_name):  # noqa: D103
+    action = get_action_query(session, action_name).first()
+    if action:
+        return action
+
+    action = Permission()
+    action.name = action_name
+    session.add(action)
+    session.commit()
+
+    return action
+
+
+def get_resource_query(session, resource_name):  # noqa: D103
+    return session.query(ViewMenu).filter(ViewMenu.name == resource_name)
+
+
+def get_action_query(session, action_name):  # noqa: D103
+    return session.query(Permission).filter(Permission.name == action_name)
+
+
+def get_pv_with_action_query(session, action):  # noqa: D103
+    return session.query(PermissionView).filter(PermissionView.permission == action)
+
+
+def get_pv_with_resource_query(session, resource):  # noqa: D103
+    return session.query(PermissionView).filter(PermissionView.view_menu_id == resource)
+
+
+def update_pv_action(session, pv_query, action):  # noqa: D103
+    pv_query.update({PermissionView.permission_id: action.id}, synchronize_session=False)
+    session.commit()
+
+
+def get_pv(session, resource, action):  # noqa: D103
+    return (
+        session.query(PermissionView)
+        .filter(PermissionView.view_menu == resource)
+        .filter(PermissionView.permission == action)
+        .first()
+    )
+
+
+def update_pv_resource(session, pv_query, resource):  # noqa: D103
+    for pv in pv_query.all():  # noqa: D103
+        if not get_pv(session, resource, pv.permission):  # noqa: D103
+            pv.view_menu = resource
+        else:
+            session.delete(pv)
+
+    session.commit()
+
+
+def migrate_to_new_dag_permissions(db):  # noqa: D103
+    # Prefix individual dag perms with `DAG:`
+    prefix_individual_dag_permissions(db.session)
+
+    # Update existing PVs to use `can_read` instead of `can_dag_read`
+    can_dag_read_action = get_action_query(db.session, 'can_dag_read').first()
+    old_can_dag_read_pvs = get_pv_with_action_query(db.session, can_dag_read_action)
+    can_read_action = get_or_create_action(db.session, 'can_read')
+    update_pv_action(db.session, old_can_dag_read_pvs, can_read_action)
+
+    # Update existing PVs to use `can_edit` instead of `can_dag_edit`
+    can_dag_edit_action = get_action_query(db.session, 'can_dag_edit').first()
+    old_can_dag_edit_pvs = get_pv_with_action_query(db.session, can_dag_edit_action)
+    can_edit_action = get_or_create_action(db.session, 'can_edit')
+    update_pv_action(db.session, old_can_dag_edit_pvs, can_edit_action)
+
+    # Update existing PVs for `all_dags` resource to use `DAGs` resource.
+    all_dags_resource = get_resource_query(db.session, 'all_dags').first()
+    if all_dags_resource:
+        old_all_dags_pv = get_pv_with_resource_query(db.session, all_dags_resource)
+        dag_resource = get_or_create_dag_resource(db.session)
+        update_pv_resource(db.session, old_all_dags_pv, dag_resource)
+
+        # Delete the `all_dags` resource
+        db.session.delete(all_dags_resource)
+
+    # Delete `can_dag_read` action
+    if can_dag_read_action:
+        db.session.delete(can_dag_read_action)
+
+    # Delete `can_dag_edit` action
+    if can_dag_edit_action:
+        db.session.delete(can_dag_edit_action)
+
+    db.session.commit()
+
+
 def upgrade():  # noqa: D103
-    permissions = ['can_dag_read', 'can_dag_edit']
-    view_menus = cached_app().appbuilder.sm.get_all_view_menu()
-    convert_permissions(permissions, view_menus, upgrade_action, upgrade_dag_id)
+    try:
+        db = SQLA()
+        db.session = settings.Session
+        migrate_to_new_dag_permissions(db)
+        db.session.commit()
+    except:  # noqa: E722

Review comment:
       Do we need this? Can't we just let it throw?

##########
File path: airflow/migrations/versions/849da589634d_prefix_dag_permissions.py
##########
@@ -34,80 +37,139 @@
 depends_on = None
 
 
+def prefix_individual_dag_permissions(session):  # noqa: D103
+    dag_perms = ['can_dag_read', 'can_dag_edit']
+    prefix = "DAG:"
+    permission_view_menus = (
+        session.query(PermissionView)
+        .join(Permission)
+        .filter(Permission.name.in_(dag_perms))
+        .join(ViewMenu)
+        .filter(ViewMenu.name != 'all_dags')
+        .filter(ViewMenu.name.notlike('DAG:%'))

Review comment:
       ```suggestion
           .filter(ViewMenu.name.notlike(prefix + '%'))
   ```
   




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org