From 52b0f5f89504e674e83e74d5117a53c6efdf06bc Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Mon, 22 Aug 2022 14:39:04 -0400
Subject: [PATCH] add ordering to grid_data

---
 airflow/www/static/js/api/useGridData.ts      | 11 +++-
 airflow/www/static/js/dag/details/Header.tsx  |  6 +-
 .../static/js/dag/grid/dagRuns/Tooltip.tsx    | 58 ++++++++++---------
 airflow/www/static/js/types/index.ts          |  3 +
 airflow/www/static/js/utils/index.test.ts     | 28 ++++++++-
 airflow/www/static/js/utils/index.ts          | 13 ++++-
 airflow/www/views.py                          |  1 +
 tests/www/views/test_views_grid.py            |  3 +
 8 files changed, 89 insertions(+), 34 deletions(-)

diff --git a/airflow/www/static/js/api/useGridData.ts b/airflow/www/static/js/api/useGridData.ts
index dc5e5dfb9a079..2e2ca1cb950c1 100644
--- a/airflow/www/static/js/api/useGridData.ts
+++ b/airflow/www/static/js/api/useGridData.ts
@@ -26,7 +26,8 @@ import useErrorToast from 'src/utils/useErrorToast';
 import useFilters, {
   BASE_DATE_PARAM, NUM_RUNS_PARAM, RUN_STATE_PARAM, RUN_TYPE_PARAM, now,
 } from 'src/dag/useFilters';
-import type { Task, DagRun } from 'src/types';
+import type { Task, DagRun, RunOrdering } from 'src/types';
+import { camelCase } from 'lodash';
 
 const DAG_ID_PARAM = 'dag_id';
 
@@ -38,6 +39,7 @@ const urlRoot = getMetaValue('root');
 interface GridData {
   dagRuns: DagRun[];
   groups: Task;
+  ordering: RunOrdering;
 }
 
 const emptyGridData: GridData = {
@@ -47,8 +49,14 @@ const emptyGridData: GridData = {
     label: null,
     instances: [],
   },
+  ordering: [],
 };
 
+const formatOrdering = (data: GridData) => ({
+  ...data,
+  ordering: data.ordering.map((o: string) => camelCase(o)) as RunOrdering,
+});
+
 export const areActiveRuns = (runs: DagRun[] = []) => runs.filter((run) => ['manual', 'manual'].includes(run.runType)).filter((run) => ['queued', 'running', 'scheduled'].includes(run.state)).length > 0;
 
 const useGridData = () => {
@@ -88,6 +96,7 @@ const useGridData = () => {
         });
         throw (error);
       },
+      select: formatOrdering,
     },
   );
   return {
diff --git a/airflow/www/static/js/dag/details/Header.tsx b/airflow/www/static/js/dag/details/Header.tsx
index 3b957f6a1779f..d7ddb98564717 100644
--- a/airflow/www/static/js/dag/details/Header.tsx
+++ b/airflow/www/static/js/dag/details/Header.tsx
@@ -25,7 +25,7 @@ import {
   Text,
 } from '@chakra-ui/react';
 
-import { getMetaValue, getTask } from 'src/utils';
+import { getDagRunLabel, getMetaValue, getTask } from 'src/utils';
 import useSelection from 'src/dag/useSelection';
 import Time from 'src/components/Time';
 import { useGridData } from 'src/api';
@@ -36,7 +36,7 @@ import BreadcrumbText from './BreadcrumbText';
 const dagId = getMetaValue('dag_id');
 
 const Header = () => {
-  const { data: { dagRuns, groups } } = useGridData();
+  const { data: { dagRuns, groups, ordering } } = useGridData();
 
   const { selected: { taskId, runId, mapIndex }, onSelect, clearSelection } = useSelection();
   const dagRun = dagRuns.find((r) => r.runId === runId);
@@ -58,7 +58,7 @@ const Header = () => {
       || runId.includes('backfill__')
       || runId.includes('dataset_triggered__')
     )
-      ? <Time dateTime={dagRun.dataIntervalStart || dagRun.executionDate} />
+      ? <Time dateTime={getDagRunLabel({ dagRun, ordering })} />
       : runId;
     runLabel = (
       <>
diff --git a/airflow/www/static/js/dag/grid/dagRuns/Tooltip.tsx b/airflow/www/static/js/dag/grid/dagRuns/Tooltip.tsx
index ae52d9bdaa24c..f8b577a1bc54c 100644
--- a/airflow/www/static/js/dag/grid/dagRuns/Tooltip.tsx
+++ b/airflow/www/static/js/dag/grid/dagRuns/Tooltip.tsx
@@ -19,9 +19,12 @@
 
 import React from 'react';
 import { Box, Text } from '@chakra-ui/react';
+import { startCase } from 'lodash';
 
 import { formatDuration } from 'src/datetime_utils';
 import Time from 'src/components/Time';
+import { getDagRunLabel } from 'src/utils';
+import { useGridData } from 'src/api';
 
 import type { RunWithDuration } from './index';
 
@@ -29,33 +32,32 @@ interface Props {
   dagRun: RunWithDuration;
 }
 
-const DagRunTooltip = ({
-  dagRun: {
-    state, duration, dataIntervalStart, executionDate, runType,
-  },
-}: Props) => (
-  <Box py="2px">
-    <Text>
-      Status:
-      {' '}
-      {state || 'no status'}
-    </Text>
-    <Text whiteSpace="nowrap">
-      Run:
-      {' '}
-      <Time dateTime={dataIntervalStart || executionDate} />
-    </Text>
-    <Text>
-      Duration:
-      {' '}
-      {formatDuration(duration)}
-    </Text>
-    <Text>
-      Type:
-      {' '}
-      {runType}
-    </Text>
-  </Box>
-);
+const DagRunTooltip = ({ dagRun }: Props) => {
+  const { data: { ordering } } = useGridData();
+  return (
+    <Box py="2px">
+      <Text>
+        Status:
+        {' '}
+        {dagRun.state || 'no status'}
+      </Text>
+      <Text whiteSpace="nowrap">
+        {startCase(ordering[0] || ordering[1])}
+        {': '}
+        <Time dateTime={getDagRunLabel({ dagRun, ordering })} />
+      </Text>
+      <Text>
+        Duration:
+        {' '}
+        {formatDuration(dagRun.duration)}
+      </Text>
+      <Text>
+        Type:
+        {' '}
+        {dagRun.runType}
+      </Text>
+    </Box>
+  );
+};
 
 export default DagRunTooltip;
diff --git a/airflow/www/static/js/types/index.ts b/airflow/www/static/js/types/index.ts
index 6918320f1563e..34f40f4dbc773 100644
--- a/airflow/www/static/js/types/index.ts
+++ b/airflow/www/static/js/types/index.ts
@@ -79,6 +79,8 @@ interface Task {
   hasOutletDatasets?: boolean;
 }
 
+type RunOrdering = ('dataIntervalStart' | 'executionDate' | 'dataIntervalEnd')[];
+
 export type {
   Dag,
   DagRun,
@@ -87,4 +89,5 @@ export type {
   TaskInstance,
   Task,
   API,
+  RunOrdering,
 };
diff --git a/airflow/www/static/js/utils/index.test.ts b/airflow/www/static/js/utils/index.test.ts
index e3d9f693866f6..03d3cf4d825f8 100644
--- a/airflow/www/static/js/utils/index.test.ts
+++ b/airflow/www/static/js/utils/index.test.ts
@@ -18,7 +18,8 @@
  */
 
 import { isEmpty } from 'lodash';
-import { getTask, getTaskSummary } from '.';
+import type { DagRun } from 'src/types';
+import { getDagRunLabel, getTask, getTaskSummary } from '.';
 
 const sampleTasks = {
   id: null,
@@ -112,3 +113,28 @@ describe('Test getTaskSummary()', () => {
     expect(isEmpty(summary.operators)).toBeTruthy();
   });
 });
+
+describe('Test getDagRunLabel', () => {
+  const dagRun = {
+    dagId: 'dagId',
+    runId: 'run1',
+    dataIntervalStart: '2021-12-07T21:14:19.704433+00:00',
+    dataIntervalEnd: '2021-12-08T21:14:19.704433+00:00',
+    startDate: '2021-11-08T21:14:19.704433+00:00',
+    endDate: '2021-11-08T21:17:13.206426+00:00',
+    state: 'failed',
+    runType: 'scheduled',
+    executionDate: '2021-12-09T21:14:19.704433+00:00',
+    lastSchedulingDecision: '2021-11-08T21:14:19.704433+00:00',
+  } as DagRun;
+
+  test('Defaults to dataIntervalStart', async () => {
+    const runLabel = getDagRunLabel({ dagRun });
+    expect(runLabel).toBe(dagRun.dataIntervalStart);
+  });
+
+  test('Passing an order overrides default', async () => {
+    const runLabel = getDagRunLabel({ dagRun, ordering: ['executionDate'] });
+    expect(runLabel).toBe(dagRun.executionDate);
+  });
+});
diff --git a/airflow/www/static/js/utils/index.ts b/airflow/www/static/js/utils/index.ts
index ffbdc0028951e..a8506eaf919f7 100644
--- a/airflow/www/static/js/utils/index.ts
+++ b/airflow/www/static/js/utils/index.ts
@@ -17,7 +17,7 @@
  * under the License.
  */
 
-import type { Task } from 'src/types';
+import type { DagRun, RunOrdering, Task } from 'src/types';
 
 // Delay in ms for various hover actions
 const hoverDelay = 200;
@@ -114,6 +114,16 @@ const getTaskSummary = ({
   };
 };
 
+interface RunLabelProps {
+  dagRun: DagRun;
+  ordering?: RunOrdering;
+}
+
+const getDagRunLabel = ({
+  dagRun,
+  ordering = ['dataIntervalEnd', 'executionDate'],
+}: RunLabelProps) => dagRun[ordering[0]] ?? dagRun[ordering[1]];
+
 export {
   hoverDelay,
   finalStatesMap,
@@ -121,4 +131,5 @@ export {
   appendSearchParams,
   getTask,
   getTaskSummary,
+  getDagRunLabel,
 };
diff --git a/airflow/www/views.py b/airflow/www/views.py
index 7febf9a31c306..a84e89547910f 100644
--- a/airflow/www/views.py
+++ b/airflow/www/views.py
@@ -3632,6 +3632,7 @@ def grid_data(self):
             data = {
                 'groups': dag_to_grid(dag, dag_runs, session),
                 'dag_runs': encoded_runs,
+                'ordering': dag.timetable.run_ordering,
             }
         # avoid spaces to reduce payload size
         return (
diff --git a/tests/www/views/test_views_grid.py b/tests/www/views/test_views_grid.py
index 6ca422c68a176..9c756ee5e607b 100644
--- a/tests/www/views/test_views_grid.py
+++ b/tests/www/views/test_views_grid.py
@@ -124,6 +124,7 @@ def test_no_runs(admin_client, dag_without_runs):
             'instances': [],
             'label': None,
         },
+        'ordering': ['data_interval_end', 'execution_date'],
     }
 
 
@@ -272,6 +273,7 @@ def test_one_run(admin_client, dag_with_runs: List[DagRun], session):
             'instances': [],
             'label': None,
         },
+        'ordering': ['data_interval_end', 'execution_date'],
     }
 
 
@@ -323,6 +325,7 @@ def _expected_task_details(task_id, has_outlet_datasets):
             'instances': [],
             'label': None,
         },
+        'ordering': ['data_interval_end', 'execution_date'],
     }