UI changes for DAG Reparsing feature (#39636)

* UI changes for DAG Reparsing feature releated: apache/airflow#39138 * Add correct permission for UI * Fix JS * Update airflow/www/static/js/main.js * Add a wrapper arounf REST DAG Reparsing API * Move file_token from dag object to a separate field * Fix CI issue * Change icons GitOrigin-RevId: 2250a1a228f30d13ae3c481b059e05bbd29d3839
GoogleCloudPlatform · Nov 9, 2024 · 7ce6970 · 7ce6970
1 parent b80a09c
commit 7ce6970
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 0 deletions.
diff --git a/airflow/www/templates/airflow/dag.html b/airflow/www/templates/airflow/dag.html
@@ -231,6 +231,12 @@ <h4 class="js-dataset-triggered" style="user-select: none;-moz-user-select: auto
             <span class="material-icons" aria-hidden="true">play_arrow</span>
         </a>
         {% endif %}
+        <a href="{{ url_for('Airflow.parse_dag', file_token=dag_file_token, redirect_url=url_for(request.endpoint, dag_id=dag.dag_id)) }}"
+           title="Reparse&nbsp;DAG"
+           aria-label="Reparse DAG"
+           class="btn btn-default btn-icon-only {{ ' disabled' if not dag.can_edit }}">
+          <span class="material-icons">restore_page</span>
+        </a>
         <a href="{{ url_for('Airflow.delete', dag_id=dag.dag_id, redirect_url=url_for(request.endpoint, dag_id=dag.dag_id)) }}"
           title="Delete&nbsp;DAG"
           aria-label="Delete DAG"

diff --git a/airflow/www/templates/airflow/dags.html b/airflow/www/templates/airflow/dags.html
@@ -392,6 +392,10 @@ <h2>{{ page_title }}</h2>
                     <span class="material-icons" aria-hidden="true">play_arrow</span>
                   </a>
                 {% endif %}
+                <a href="{{ url_for('Airflow.parse_dag', file_token=file_tokens[dag.dag_id], redirect_url=url_for(request.endpoint)) }}" title="Reparse&nbsp;DAG" aria-label="Reparse DAG"
+                    class="btn btn-sm btn-default btn-icon-only{{ ' disabled' if not dag.can_edit }}">
+                    <span class="material-icons">restore_page</span>
+                </a>
                 <a href="{{ url_for('Airflow.delete', dag_id=dag.dag_id, redirect_url=url_for(request.endpoint)) }}"
                     onclick="return confirmDeleteDag(this, '{{ dag.dag_id }}')" title="Delete&nbsp;DAG"
                     aria-label="Delete DAG"

diff --git a/airflow/www/views.py b/airflow/www/views.py
@@ -48,6 +48,7 @@
     Response,
     abort,
     before_render_template,
+    current_app,
     flash,
     g,
     has_request_context,
@@ -67,6 +68,7 @@
 from flask_appbuilder.urltools import get_order_args, get_page_args, get_page_size_args
 from flask_appbuilder.widgets import FormWidget
 from flask_babel import lazy_gettext
+from itsdangerous import URLSafeSerializer
 from jinja2.utils import htmlsafe_json_dumps, pformat  # type: ignore
 from markupsafe import Markup, escape
 from pendulum.datetime import DateTime
@@ -739,6 +741,8 @@ def render_template(self, *args, **kwargs):
             kwargs["can_edit_dag"] = get_auth_manager().is_authorized_dag(
                 method="PUT", details=DagDetails(id=kwargs["dag"].dag_id)
             )
+            url_serializer = URLSafeSerializer(current_app.config["SECRET_KEY"])
+            kwargs["dag_file_token"] = url_serializer.dumps(kwargs["dag"].fileloc)
 
         return super().render_template(
             *args,
@@ -957,6 +961,7 @@ def index(self):
             else:
                 dataset_triggered_next_run_info = {}
 
+            file_tokens = {}
             for dag in dags:
                 dag.can_edit = get_auth_manager().is_authorized_dag(
                     method="PUT", details=DagDetails(id=dag.dag_id), user=g.user
@@ -965,6 +970,8 @@ def index(self):
                 dag.can_delete = get_auth_manager().is_authorized_dag(
                     method="DELETE", details=DagDetails(id=dag.dag_id), user=g.user
                 )
+                url_serializer = URLSafeSerializer(current_app.config["SECRET_KEY"])
+                file_tokens[dag.dag_id] = url_serializer.dumps(dag.fileloc)
 
             dagtags = session.execute(select(func.distinct(DagTag.name)).order_by(DagTag.name)).all()
             tags = [
@@ -1106,6 +1113,7 @@ def _iter_parsed_moved_data_table_names():
             auto_refresh_interval=conf.getint("webserver", "auto_refresh_interval"),
             dataset_triggered_next_run_info=dataset_triggered_next_run_info,
             scarf_url=scarf_url,
+            file_tokens=file_tokens,
         )
 
     @expose("/datasets")
@@ -1404,6 +1412,7 @@ def rendered_templates(self, session):
         raw_task = dag.get_task(task_id).prepare_for_execution()
 
         no_dagrun = False
+        url_serializer = URLSafeSerializer(current_app.config["SECRET_KEY"])
 
         title = "Rendered Template"
         html_dict = {}
@@ -1504,6 +1513,7 @@ def rendered_templates(self, session):
             form=form,
             root=root,
             title=title,
+            dag_file_token=url_serializer.dumps(dag.fileloc),
         )
 
     @expose("/rendered-k8s")
@@ -2880,6 +2890,7 @@ def grid(self, dag_id: str, session: Session = NEW_SESSION):
         color_log_warning_keywords = conf.get("logging", "color_log_warning_keywords", fallback="")
 
         dag = get_airflow_app().dag_bag.get_dag(dag_id, session=session)
+        url_serializer = URLSafeSerializer(current_app.config["SECRET_KEY"])
         dag_model = DagModel.get_dagmodel(dag_id, session=session)
         if not dag:
             flash(f'DAG "{dag_id}" seems to be missing from DagBag.', "error")
@@ -2938,6 +2949,7 @@ def grid(self, dag_id: str, session: Session = NEW_SESSION):
             excluded_events_raw=excluded_events_raw,
             color_log_error_keywords=color_log_error_keywords,
             color_log_warning_keywords=color_log_warning_keywords,
+            dag_file_token=url_serializer.dumps(dag.fileloc),
         )
 
     @expose("/calendar")
@@ -3591,6 +3603,22 @@ def audit_log(self, dag_id: str):
 
         return redirect(url_for("Airflow.grid", **kwargs))
 
+    @expose("/parseDagFile/<string:file_token>")
+    def parse_dag(self, file_token: str):
+        from airflow.api_connexion.endpoints.dag_parsing import reparse_dag_file
+
+        with create_session() as session:
+            response = reparse_dag_file(file_token=file_token, session=session)
+            response_messages = {
+                201: ["Reparsing request submitted successfully", "info"],
+                401: ["Unauthenticated request", "error"],
+                403: ["Permission Denied", "error"],
+                404: ["DAG not found", "error"],
+            }
+            flash(response_messages[response.status_code][0], response_messages[response.status_code][1])
+        redirect_url = get_safe_url(request.values.get("redirect_url"))
+        return redirect(redirect_url)
+
 
 class ConfigurationView(AirflowBaseView):
     """View to show Airflow Configurations."""