From 0aae6bbb29d7362a288706583cfbe788fbf8673d Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Thu, 2 Jun 2022 19:48:53 +0200
Subject: [PATCH] Reduce grid view API calls (#24083)

* Reduce API calls from /grid

- Separate /grid_data from /grid
- Remove need for formatData
- Increase default query stale time to prevent extra fetches
- Fix useTask query keys

* consolidate grid data functions

* fix www tests

test grid_data instead of /grid

(cherry picked from commit 035553c86988f403b43ef5825715b45f055d62dd)
---
 airflow/www/static/js/grid/Main.jsx           | 26 +++++--
 airflow/www/static/js/grid/api/useGridData.js |  7 +-
 .../useGridData.test.js}                      |  2 +-
 .../static/js/grid/api/useGridData.test.jsx   | 78 -------------------
 airflow/www/static/js/grid/api/useTasks.js    |  4 +-
 .../static/js/grid/context/autorefresh.jsx    | 15 +---
 .../www/static/js/grid/dagRuns/index.test.jsx | 10 +--
 .../static/js/grid/details/content/Dag.jsx    |  3 +-
 airflow/www/static/js/grid/index.jsx          |  1 +
 airflow/www/static/js/grid/utils/gridData.js  | 34 --------
 airflow/www/templates/airflow/grid.html       |  1 -
 airflow/www/views.py                          | 11 ---
 tests/www/views/test_views_acl.py             | 10 +--
 tests/www/views/test_views_tasks.py           | 18 +----
 14 files changed, 38 insertions(+), 182 deletions(-)
 rename airflow/www/static/js/grid/{utils/gridData.test.js => api/useGridData.test.js} (96%)
 delete mode 100644 airflow/www/static/js/grid/api/useGridData.test.jsx
 delete mode 100644 airflow/www/static/js/grid/utils/gridData.js

diff --git a/airflow/www/static/js/grid/Main.jsx b/airflow/www/static/js/grid/Main.jsx
index 98aa360aa6789..5e9a4f2a9a12d 100644
--- a/airflow/www/static/js/grid/Main.jsx
+++ b/airflow/www/static/js/grid/Main.jsx
@@ -25,17 +25,21 @@ import {
   Flex,
   useDisclosure,
   Divider,
+  Spinner,
 } from '@chakra-ui/react';
+import { isEmpty } from 'lodash';
 
 import Details from './details';
 import useSelection from './utils/useSelection';
 import Grid from './Grid';
 import FilterBar from './FilterBar';
 import LegendRow from './LegendRow';
+import { useGridData } from './api';
 
 const detailsPanelKey = 'hideDetailsPanel';
 
 const Main = () => {
+  const { data: { groups }, isLoading } = useGridData();
   const isPanelOpen = localStorage.getItem(detailsPanelKey) !== 'true';
   const { isOpen, onToggle } = useDisclosure({ defaultIsOpen: isPanelOpen });
   const { clearSelection } = useSelection();
@@ -57,14 +61,20 @@ const Main = () => {
       <LegendRow setHoveredTaskState={setHoveredTaskState} />
       <Divider mb={5} borderBottomWidth={2} />
       <Flex justifyContent="space-between">
-        <Grid
-          isPanelOpen={isOpen}
-          onPanelToggle={onPanelToggle}
-          hoveredTaskState={hoveredTaskState}
-        />
-        <Box borderLeftWidth={isOpen ? 1 : 0} position="relative">
-          {isOpen && (<Details />)}
-        </Box>
+        {isLoading || isEmpty(groups)
+          ? (<Spinner />)
+          : (
+            <>
+              <Grid
+                isPanelOpen={isOpen}
+                onPanelToggle={onPanelToggle}
+                hoveredTaskState={hoveredTaskState}
+              />
+              <Box borderLeftWidth={isOpen ? 1 : 0} position="relative">
+                {isOpen && (<Details />)}
+              </Box>
+            </>
+          )}
       </Flex>
     </Box>
   );
diff --git a/airflow/www/static/js/grid/api/useGridData.js b/airflow/www/static/js/grid/api/useGridData.js
index d31712989bcfd..38d4e00748d32 100644
--- a/airflow/www/static/js/grid/api/useGridData.js
+++ b/airflow/www/static/js/grid/api/useGridData.js
@@ -17,14 +17,13 @@
  * under the License.
  */
 
-/* global autoRefreshInterval, gridData */
+/* global autoRefreshInterval */
 
 import { useQuery } from 'react-query';
 import axios from 'axios';
 
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import { areActiveRuns, formatData } from '../utils/gridData';
 import useErrorToast from '../utils/useErrorToast';
 import useFilters, {
   BASE_DATE_PARAM, NUM_RUNS_PARAM, RUN_STATE_PARAM, RUN_TYPE_PARAM, now,
@@ -42,8 +41,9 @@ const emptyData = {
   groups: {},
 };
 
+export const areActiveRuns = (runs = []) => runs.filter((run) => ['queued', 'running', 'scheduled'].includes(run.state)).length > 0;
+
 const useGridData = () => {
-  const initialData = formatData(gridData, emptyData);
   const { isRefreshOn, stopRefresh } = useAutoRefresh();
   const errorToast = useErrorToast();
   const {
@@ -75,7 +75,6 @@ const useGridData = () => {
       throw (error);
     }
   }, {
-    initialData,
     placeholderData: emptyData,
     // only refetch if the refresh switch is on
     refetchInterval: isRefreshOn && autoRefreshInterval * 1000,
diff --git a/airflow/www/static/js/grid/utils/gridData.test.js b/airflow/www/static/js/grid/api/useGridData.test.js
similarity index 96%
rename from airflow/www/static/js/grid/utils/gridData.test.js
rename to airflow/www/static/js/grid/api/useGridData.test.js
index 6bbd8bf8b6a6d..29a7f1ac8a11f 100644
--- a/airflow/www/static/js/grid/utils/gridData.test.js
+++ b/airflow/www/static/js/grid/api/useGridData.test.js
@@ -19,7 +19,7 @@
 
 /* global describe, test, expect */
 
-import { areActiveRuns } from './gridData';
+import { areActiveRuns } from './useGridData';
 
 describe('Test areActiveRuns()', () => {
   test('Correctly detects active runs', () => {
diff --git a/airflow/www/static/js/grid/api/useGridData.test.jsx b/airflow/www/static/js/grid/api/useGridData.test.jsx
deleted file mode 100644
index 24aece6f5919f..0000000000000
--- a/airflow/www/static/js/grid/api/useGridData.test.jsx
+++ /dev/null
@@ -1,78 +0,0 @@
-/*!
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-/* global describe, test, expect, beforeAll */
-
-import { renderHook } from '@testing-library/react-hooks';
-import useGridData from './useGridData';
-import { Wrapper } from '../utils/testUtils';
-
-const pendingGridData = {
-  groups: {},
-  dag_runs: [
-    {
-      dag_id: 'example_python_operator',
-      run_id: 'manual__2021-11-08T21:14:17.170046+00:00',
-      start_date: null,
-      end_date: null,
-      state: 'queued',
-      execution_date: '2021-11-08T21:14:17.170046+00:00',
-      data_interval_start: '2021-11-08T21:14:17.170046+00:00',
-      data_interval_end: '2021-11-08T21:14:17.170046+00:00',
-      run_type: 'manual',
-    },
-  ],
-};
-
-describe('Test useGridData hook', () => {
-  beforeAll(() => {
-    global.autoRefreshInterval = 5;
-  });
-
-  test('data is valid camelcase json', () => {
-    global.gridData = JSON.stringify(pendingGridData);
-
-    const { result } = renderHook(() => useGridData(), { wrapper: Wrapper });
-    const { data } = result.current;
-
-    expect(typeof data === 'object').toBe(true);
-    expect(data.dagRuns).toBeDefined();
-    expect(data.dag_runs).toBeUndefined();
-  });
-
-  test('Can handle no gridData', () => {
-    global.gridData = null;
-
-    const { result } = renderHook(() => useGridData(), { wrapper: Wrapper });
-    const { data } = result.current;
-
-    expect(data.dagRuns).toStrictEqual([]);
-    expect(data.groups).toStrictEqual({});
-  });
-
-  test('Can handle empty gridData object', () => {
-    global.gridData = {};
-
-    const { result } = renderHook(() => useGridData(), { wrapper: Wrapper });
-    const { data } = result.current;
-
-    expect(data.dagRuns).toStrictEqual([]);
-    expect(data.groups).toStrictEqual({});
-  });
-});
diff --git a/airflow/www/static/js/grid/api/useTasks.js b/airflow/www/static/js/grid/api/useTasks.js
index 9d0f56d7cb59c..a635622e2787b 100644
--- a/airflow/www/static/js/grid/api/useTasks.js
+++ b/airflow/www/static/js/grid/api/useTasks.js
@@ -23,9 +23,9 @@ import { getMetaValue } from '../../utils';
 
 const tasksUrl = getMetaValue('tasks_api');
 
-export default function useTasks(dagId) {
+export default function useTasks() {
   return useQuery(
-    ['tasks', dagId],
+    'tasks',
     () => axios.get(tasksUrl),
     {
       placeholderData: { tasks: [] },
diff --git a/airflow/www/static/js/grid/context/autorefresh.jsx b/airflow/www/static/js/grid/context/autorefresh.jsx
index ae242e5bf66aa..7fbdc405f93c1 100644
--- a/airflow/www/static/js/grid/context/autorefresh.jsx
+++ b/airflow/www/static/js/grid/context/autorefresh.jsx
@@ -17,11 +17,10 @@
  * under the License.
  */
 
-/* global localStorage, gridData, document */
+/* global localStorage, document */
 
 import React, { useContext, useState, useEffect } from 'react';
 import { getMetaValue } from '../../utils';
-import { formatData, areActiveRuns } from '../utils/gridData';
 
 const autoRefreshKey = 'disabledAutoRefresh';
 
@@ -31,17 +30,9 @@ const isRefreshDisabled = JSON.parse(localStorage.getItem(autoRefreshKey));
 const AutoRefreshContext = React.createContext(null);
 
 export const AutoRefreshProvider = ({ children }) => {
-  let dagRuns = [];
-  try {
-    const data = JSON.parse(gridData);
-    if (data.dag_runs) dagRuns = formatData(data.dag_runs);
-  } catch {
-    dagRuns = [];
-  }
   const [isPaused, setIsPaused] = useState(initialIsPaused);
-  const isActive = areActiveRuns(dagRuns);
   const isRefreshAllowed = !(isPaused || isRefreshDisabled);
-  const initialState = isRefreshAllowed && isActive;
+  const initialState = isRefreshAllowed;
 
   const [isRefreshOn, setRefresh] = useState(initialState);
 
@@ -67,8 +58,6 @@ export const AutoRefreshProvider = ({ children }) => {
       setIsPaused(!e.value);
       if (!e.value) {
         stopRefresh();
-      } else if (isActive) {
-        setRefresh(true);
       }
     };
 
diff --git a/airflow/www/static/js/grid/dagRuns/index.test.jsx b/airflow/www/static/js/grid/dagRuns/index.test.jsx
index a507df522c565..0420d7aba8792 100644
--- a/airflow/www/static/js/grid/dagRuns/index.test.jsx
+++ b/airflow/www/static/js/grid/dagRuns/index.test.jsx
@@ -112,15 +112,7 @@ describe('Test DagRuns', () => {
   });
 
   test('Handles empty data correctly', () => {
-    global.gridData = null;
-    const { queryByTestId } = render(
-      <DagRuns />, { wrapper: TableWrapper },
-    );
-    expect(queryByTestId('run')).toBeNull();
-  });
-
-  test('Handles no data correctly', () => {
-    global.gridData = {};
+    global.autoRefreshInterval = 0;
     const { queryByTestId } = render(
       <DagRuns />, { wrapper: TableWrapper },
     );
diff --git a/airflow/www/static/js/grid/details/content/Dag.jsx b/airflow/www/static/js/grid/details/content/Dag.jsx
index 0f434d88fd72f..e527b494cf283 100644
--- a/airflow/www/static/js/grid/details/content/Dag.jsx
+++ b/airflow/www/static/js/grid/details/content/Dag.jsx
@@ -37,11 +37,10 @@ import { useTasks, useGridData } from '../../api';
 import Time from '../../components/Time';
 import { SimpleStatus } from '../../components/StatusBox';
 
-const dagId = getMetaValue('dag_id');
 const dagDetailsUrl = getMetaValue('dag_details_url');
 
 const Dag = () => {
-  const { data: taskData } = useTasks(dagId);
+  const { data: taskData } = useTasks();
   const { data: { dagRuns } } = useGridData();
   if (!taskData) return null;
   const { tasks = [], totalEntries = '' } = taskData;
diff --git a/airflow/www/static/js/grid/index.jsx b/airflow/www/static/js/grid/index.jsx
index 3d8702841edc6..b8c6db5324d4b 100644
--- a/airflow/www/static/js/grid/index.jsx
+++ b/airflow/www/static/js/grid/index.jsx
@@ -48,6 +48,7 @@ const queryClient = new QueryClient({
       refetchOnWindowFocus: false,
       retry: 1,
       retryDelay: 500,
+      staleTime: 5 * 60 * 1000, // 5 minutes
       refetchOnMount: true, // Refetches stale queries, not "always"
     },
     mutations: {
diff --git a/airflow/www/static/js/grid/utils/gridData.js b/airflow/www/static/js/grid/utils/gridData.js
deleted file mode 100644
index 95171652b7a7b..0000000000000
--- a/airflow/www/static/js/grid/utils/gridData.js
+++ /dev/null
@@ -1,34 +0,0 @@
-/*!
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements.  See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership.  The ASF licenses this file
- * to you under the Apache License, Version 2.0 (the
- * "License"); you may not use this file except in compliance
- * with the License.  You may obtain a copy of the License at
- *
- *   http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing,
- * software distributed under the License is distributed on an
- * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
- * KIND, either express or implied.  See the License for the
- * specific language governing permissions and limitations
- * under the License.
- */
-
-import camelcaseKeys from 'camelcase-keys';
-
-export const areActiveRuns = (runs = []) => runs.filter((run) => ['queued', 'running', 'scheduled'].includes(run.state)).length > 0;
-
-export const formatData = (data, emptyData) => {
-  if (!data || !Object.keys(data).length) {
-    return emptyData;
-  }
-  let formattedData = data;
-  // Convert to json if needed
-  if (typeof data === 'string') formattedData = JSON.parse(data);
-  // change from pascal to camelcase
-  formattedData = camelcaseKeys(formattedData, { deep: true });
-  return formattedData;
-};
diff --git a/airflow/www/templates/airflow/grid.html b/airflow/www/templates/airflow/grid.html
index 46d52f7d27f58..f72e0ed4f1704 100644
--- a/airflow/www/templates/airflow/grid.html
+++ b/airflow/www/templates/airflow/grid.html
@@ -39,7 +39,6 @@
 {% block tail_js %}
   {{ super()}}
   <script>
-    const gridData = {{ data|tojson }};
     const stateColors = {{ state_color_mapping|tojson }};
     const autoRefreshInterval = {{ auto_refresh_interval }};
     const defaultDagRunDisplayNumber = {{ default_dag_run_display_number }};
diff --git a/airflow/www/views.py b/airflow/www/views.py
index fbebef3ecf5d6..2c306cd9445b7 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -2612,8 +2612,6 @@ def grid(self, dag_id, session=None):
             .limit(num_runs)
             .all()
         )
-        dag_runs.reverse()
-        encoded_runs = [wwwutils.encode_dag_run(dr) for dr in dag_runs]
         dag_run_dates = {dr.execution_date: alchemy_to_dict(dr) for dr in dag_runs}
 
         max_date = max(dag_run_dates, default=None)
@@ -2633,14 +2631,6 @@ def grid(self, dag_id, session=None):
         else:
             external_log_name = None
 
-        data = {
-            'groups': task_group_to_grid(dag.task_group, dag, dag_runs, session),
-            'dag_runs': encoded_runs,
-        }
-
-        # avoid spaces to reduce payload size
-        data = htmlsafe_json_dumps(data, separators=(',', ':'))
-
         default_dag_run_display_number = conf.getint('webserver', 'default_dag_run_display_number')
 
         num_runs_options = [5, 25, 50, 100, 365]
@@ -2655,7 +2645,6 @@ def grid(self, dag_id, session=None):
             form=form,
             dag=dag,
             doc_md=doc_md,
-            data=data,
             num_runs=num_runs,
             show_external_log_redirect=task_log_reader.supports_external_link,
             external_log_name=external_log_name,
diff --git a/tests/www/views/test_views_acl.py b/tests/www/views/test_views_acl.py
index d96098a5c6d5d..84003bf366017 100644
--- a/tests/www/views/test_views_acl.py
+++ b/tests/www/views/test_views_acl.py
@@ -563,7 +563,7 @@ def client_dags_tis_logs(acl_app, user_dags_tis_logs):
 TRIES_URL = "tries?days=30&dag_id=example_bash_operator"
 LANDING_TIMES_URL = "landing_times?days=30&dag_id=example_bash_operator"
 GANTT_URL = "gantt?dag_id=example_bash_operator"
-TREE_URL = "tree?dag_id=example_bash_operator"
+GRID_DATA_URL = "object/grid_data?dag_id=example_bash_operator"
 LOG_URL = (
     f"log?task_id=runme_0&dag_id=example_bash_operator&"
     f"execution_date={urllib.parse.quote_plus(str(DEFAULT_DATE))}"
@@ -581,8 +581,8 @@ def client_dags_tis_logs(acl_app, user_dags_tis_logs):
         ("client_all_dags_tis", TRIES_URL, "example_bash_operator"),
         ("client_all_dags_tis", LANDING_TIMES_URL, "example_bash_operator"),
         ("client_all_dags_tis", GANTT_URL, "example_bash_operator"),
-        ("client_dags_tis_logs", TREE_URL, "runme_1"),
-        ("viewer_client", TREE_URL, "runme_1"),
+        ("client_dags_tis_logs", GRID_DATA_URL, "runme_1"),
+        ("viewer_client", GRID_DATA_URL, "runme_1"),
         ("client_dags_tis_logs", LOG_URL, "Log by attempts"),
         ("user_client", LOG_URL, "Log by attempts"),
     ],
@@ -595,8 +595,8 @@ def client_dags_tis_logs(acl_app, user_dags_tis_logs):
         "tries",
         "landing-times",
         "gantt",
-        "tree-for-readonly-role",
-        "tree-for-viewer",
+        "grid-data-for-readonly-role",
+        "grid-data-for-viewer",
         "log",
         "log-for-user",
     ],
diff --git a/tests/www/views/test_views_tasks.py b/tests/www/views/test_views_tasks.py
index 3428ac848aecd..d99a8fd5cacdc 100644
--- a/tests/www/views/test_views_tasks.py
+++ b/tests/www/views/test_views_tasks.py
@@ -165,24 +165,14 @@ def client_ti_without_dag_edit(app):
             id='graph',
         ),
         pytest.param(
-            'tree?dag_id=example_bash_operator',
+            'object/grid_data?dag_id=example_bash_operator',
             ['runme_1'],
-            id='tree',
+            id='grid-data',
         ),
         pytest.param(
-            'dags/example_bash_operator/grid',
-            ['runme_1'],
-            id='grid',
-        ),
-        pytest.param(
-            'tree?dag_id=example_subdag_operator.section-1',
-            ['section-1-task-1'],
-            id="tree-subdag-url-param",
-        ),
-        pytest.param(
-            'dags/example_subdag_operator.section-1/grid',
+            'object/grid_data?dag_id=example_subdag_operator.section-1',
             ['section-1-task-1'],
-            id="grid-subdag",
+            id="grid-data-subdag",
         ),
         pytest.param(
             'duration?days=30&dag_id=example_bash_operator',