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
apache · Jun 2, 2022 · 035553c · 035553c
1 parent 1c680e4
commit 035553c
Show file tree

Hide file tree

Showing 14 changed files with 38 additions and 182 deletions.
diff --git 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
@@ -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/...www/static/js/grid/utils/gridData.test.js → ...ww/static/js/grid/api/useGridData.test.js b/...www/static/js/grid/utils/gridData.test.js → ...ww/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
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
diff --git 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
@@ -2611,8 +2611,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)
@@ -2632,14 +2630,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]
@@ -2654,7 +2644,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
@@ -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
@@ -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',