From 74cd7e56b4607011e34de152381ab0009070d901 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Wed, 27 Apr 2022 11:07:54 -0400
Subject: [PATCH 1/5] Add url params for dag_run_id and task_id

---
 airflow/www/package.json                      |  1 +
 airflow/www/static/js/tree/api/useTasks.js    |  3 ++
 .../www/static/js/tree/context/selection.jsx  | 35 +++++++++++++++++--
 airflow/www/static/js/tree/index.jsx          |  9 +++--
 airflow/www/yarn.lock                         | 29 +++++++++++++++
 5 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/airflow/www/package.json b/airflow/www/package.json
index 56ae182eb1ef9..bf9cc10dbe75e 100644
--- a/airflow/www/package.json
+++ b/airflow/www/package.json
@@ -95,6 +95,7 @@
     "react-dom": "^17.0.2",
     "react-icons": "^4.3.1",
     "react-query": "^3.34.16",
+    "react-router-dom": "^6.3.0",
     "react-table": "^7.7.0",
     "redoc": "^2.0.0-rc.63",
     "url-search-params-polyfill": "^8.1.0"
diff --git a/airflow/www/static/js/tree/api/useTasks.js b/airflow/www/static/js/tree/api/useTasks.js
index 86192ece5c6a8..9d0f56d7cb59c 100644
--- a/airflow/www/static/js/tree/api/useTasks.js
+++ b/airflow/www/static/js/tree/api/useTasks.js
@@ -27,5 +27,8 @@ export default function useTasks(dagId) {
   return useQuery(
     ['tasks', dagId],
     () => axios.get(tasksUrl),
+    {
+      placeholderData: { tasks: [] },
+    },
   );
 }
diff --git a/airflow/www/static/js/tree/context/selection.jsx b/airflow/www/static/js/tree/context/selection.jsx
index 29fcece22370a..289180f9a04c8 100644
--- a/airflow/www/static/js/tree/context/selection.jsx
+++ b/airflow/www/static/js/tree/context/selection.jsx
@@ -17,12 +17,17 @@
  * under the License.
  */
 
-import React, { useContext, useReducer } from 'react';
+import React, {
+  useContext, useReducer, useEffect, useRef,
+} from 'react';
+import { useSearchParams } from 'react-router-dom';
 
 const SelectionContext = React.createContext(null);
 
 const SELECT = 'SELECT';
 const DESELECT = 'DESELECT';
+const RUN_ID = 'dag_run_id';
+const TASK_ID = 'task_id';
 
 const selectionReducer = (state, { type, payload }) => {
   switch (type) {
@@ -41,10 +46,34 @@ const selectionReducer = (state, { type, payload }) => {
 
 // Expose the grid selection to any react component instead of passing around lots of props
 export const SelectionProvider = ({ children }) => {
+  const [searchParams, setSearchParams] = useSearchParams();
   const [selected, dispatch] = useReducer(selectionReducer, {});
 
-  const clearSelection = () => dispatch({ type: DESELECT });
-  const onSelect = (payload) => dispatch({ type: SELECT, payload });
+  const rendered = useRef(false);
+
+  const clearSelection = () => {
+    setSearchParams();
+    dispatch({ type: DESELECT });
+  };
+
+  const onSelect = (payload) => {
+    const params = new URLSearchParams(searchParams);
+    if (payload.runId) params.set(RUN_ID, payload.runId);
+    if (payload.taskId) params.set(TASK_ID, payload.taskId);
+    setSearchParams(params);
+    dispatch({ type: SELECT, payload });
+  };
+
+  useEffect(() => {
+    if (!rendered.current) {
+      const runId = searchParams.get(RUN_ID);
+      const taskId = searchParams.get(TASK_ID);
+      if (runId) {
+        dispatch({ type: SELECT, payload: { runId, taskId } });
+      }
+      rendered.current = true;
+    }
+  }, [searchParams]);
 
   return (
     <SelectionContext.Provider value={{ selected, clearSelection, onSelect }}>
diff --git a/airflow/www/static/js/tree/index.jsx b/airflow/www/static/js/tree/index.jsx
index 4db274f117281..5deca3a1c6d40 100644
--- a/airflow/www/static/js/tree/index.jsx
+++ b/airflow/www/static/js/tree/index.jsx
@@ -21,6 +21,7 @@
 
 import React from 'react';
 import ReactDOM from 'react-dom';
+import { BrowserRouter } from 'react-router-dom';
 import { ChakraProvider, extendTheme } from '@chakra-ui/react';
 import { CacheProvider } from '@emotion/react';
 import createCache from '@emotion/cache';
@@ -77,9 +78,11 @@ function App() {
             <QueryClientProvider client={queryClient}>
               <TimezoneProvider>
                 <AutoRefreshProvider>
-                  <SelectionProvider>
-                    <Tree />
-                  </SelectionProvider>
+                  <BrowserRouter>
+                    <SelectionProvider>
+                      <Tree />
+                    </SelectionProvider>
+                  </BrowserRouter>
                 </AutoRefreshProvider>
               </TimezoneProvider>
             </QueryClientProvider>
diff --git a/airflow/www/yarn.lock b/airflow/www/yarn.lock
index 04bc350e312eb..f342f70deb978 100644
--- a/airflow/www/yarn.lock
+++ b/airflow/www/yarn.lock
@@ -1298,6 +1298,13 @@
   dependencies:
     regenerator-runtime "^0.13.4"
 
+"@babel/runtime@^7.7.6":
+  version "7.17.9"
+  resolved "https://registry.yarnpkg.com/@babel/runtime/-/runtime-7.17.9.tgz#d19fbf802d01a8cb6cf053a64e472d42c434ba72"
+  integrity sha512-lSiBBvodq29uShpWGNbgFdKYNiFDo5/HIYsaCEY9ff4sb10x9jizo2+pRrSyF4jKZCXqgzuqBOQKbUm90gQwJg==
+  dependencies:
+    regenerator-runtime "^0.13.4"
+
 "@babel/template@^7.14.5":
   version "7.14.5"
   resolved "https://registry.yarnpkg.com/@babel/template/-/template-7.14.5.tgz#a9bc9d8b33354ff6e55a9c60d1109200a68974f4"
@@ -6473,6 +6480,13 @@ hey-listen@^1.0.8:
   resolved "https://registry.yarnpkg.com/hey-listen/-/hey-listen-1.0.8.tgz#8e59561ff724908de1aa924ed6ecc84a56a9aa68"
   integrity sha512-COpmrF2NOg4TBWUJ5UVyaCU2A88wEMkUPK4hNqyCkqHbxT92BbvfjoSozkAIIm6XhicGlJHhFdullInrdhwU8Q==
 
+history@^5.2.0:
+  version "5.3.0"
+  resolved "https://registry.yarnpkg.com/history/-/history-5.3.0.tgz#1548abaa245ba47992f063a0783db91ef201c73b"
+  integrity sha512-ZqaKwjjrAYUYfLG+htGaIIZ4nioX2L70ZUMIFysS3xvBsSG4x/n1V6TXV3N8ZYNuFGlDirFg32T7B6WOUPDYcQ==
+  dependencies:
+    "@babel/runtime" "^7.7.6"
+
 hmac-drbg@^1.0.1:
   version "1.0.1"
   resolved "https://registry.yarnpkg.com/hmac-drbg/-/hmac-drbg-1.0.1.tgz#d2745701025a6c775a6c545793ed502fc0c649a1"
@@ -9829,6 +9843,21 @@ react-remove-scroll@2.4.1:
     use-callback-ref "^1.2.3"
     use-sidecar "^1.0.1"
 
+react-router-dom@^6.3.0:
+  version "6.3.0"
+  resolved "https://registry.yarnpkg.com/react-router-dom/-/react-router-dom-6.3.0.tgz#a0216da813454e521905b5fa55e0e5176123f43d"
+  integrity sha512-uaJj7LKytRxZNQV8+RbzJWnJ8K2nPsOOEuX7aQstlMZKQT0164C+X2w6bnkqU3sjtLvpd5ojrezAyfZ1+0sStw==
+  dependencies:
+    history "^5.2.0"
+    react-router "6.3.0"
+
+react-router@6.3.0:
+  version "6.3.0"
+  resolved "https://registry.yarnpkg.com/react-router/-/react-router-6.3.0.tgz#3970cc64b4cb4eae0c1ea5203a80334fdd175557"
+  integrity sha512-7Wh1DzVQ+tlFjkeo+ujvjSqSJmkt1+8JO+T5xklPlgrh70y7ogx75ODRW0ThWhY7S+6yEDks8TYrtQe/aoboBQ==
+  dependencies:
+    history "^5.2.0"
+
 react-style-singleton@^2.1.0:
   version "2.1.1"
   resolved "https://registry.yarnpkg.com/react-style-singleton/-/react-style-singleton-2.1.1.tgz#ce7f90b67618be2b6b94902a30aaea152ce52e66"

From 8f33b7827942f5134e7c8bd1574bf464022639e9 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Wed, 27 Apr 2022 11:18:32 -0400
Subject: [PATCH 2/5] Persist other search params

---
 airflow/www/static/js/tree/context/selection.jsx | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/airflow/www/static/js/tree/context/selection.jsx b/airflow/www/static/js/tree/context/selection.jsx
index 289180f9a04c8..148ded29585f4 100644
--- a/airflow/www/static/js/tree/context/selection.jsx
+++ b/airflow/www/static/js/tree/context/selection.jsx
@@ -51,8 +51,11 @@ export const SelectionProvider = ({ children }) => {
 
   const rendered = useRef(false);
 
+  // Clear selection, but keep other search params
   const clearSelection = () => {
-    setSearchParams();
+    searchParams.delete(RUN_ID);
+    searchParams.delete(TASK_ID);
+    setSearchParams(searchParams);
     dispatch({ type: DESELECT });
   };
 
@@ -64,6 +67,7 @@ export const SelectionProvider = ({ children }) => {
     dispatch({ type: SELECT, payload });
   };
 
+  // Check search params and set selection but only on the first render
   useEffect(() => {
     if (!rendered.current) {
       const runId = searchParams.get(RUN_ID);

From 9370e01b492c8961e099c45159724f0ec215d309 Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Wed, 27 Apr 2022 12:20:08 -0400
Subject: [PATCH 3/5] simplify useSelection

---
 airflow/api_connexion/schemas/task_schema.py  |  1 +
 airflow/www/static/js/tree/Tree.jsx           |  2 +-
 airflow/www/static/js/tree/api/useClearRun.js |  2 +-
 .../www/static/js/tree/api/useClearTask.js    |  2 +-
 .../static/js/tree/api/useConfirmMarkTask.js  |  2 +-
 .../static/js/tree/api/useMarkFailedRun.js    |  2 +-
 .../static/js/tree/api/useMarkFailedTask.js   |  2 +-
 .../static/js/tree/api/useMarkSuccessRun.js   |  2 +-
 .../static/js/tree/api/useMarkSuccessTask.js  |  2 +-
 airflow/www/static/js/tree/api/useQueueRun.js |  2 +-
 airflow/www/static/js/tree/api/useRunTask.js  |  2 +-
 airflow/www/static/js/tree/api/useTreeData.js |  4 +-
 .../static/js/tree/context/autorefresh.jsx    |  2 +-
 airflow/www/static/js/tree/dagRuns/index.jsx  |  2 +-
 .../www/static/js/tree/dagRuns/index.test.jsx |  6 +--
 airflow/www/static/js/tree/details/Header.jsx |  8 +--
 airflow/www/static/js/tree/details/index.jsx  |  2 +-
 airflow/www/static/js/tree/index.jsx          |  5 +-
 airflow/www/static/js/tree/renderTaskRows.jsx |  2 +-
 .../static/js/tree/renderTaskRows.test.jsx    |  6 +--
 .../{treeDataUtils.js => utils/treeData.js}   |  0
 .../js/tree/{ => utils}/useErrorToast.js      |  0
 .../selection.jsx => utils/useSelection.jsx}  | 52 +++----------------
 23 files changed, 35 insertions(+), 75 deletions(-)
 rename airflow/www/static/js/tree/{treeDataUtils.js => utils/treeData.js} (100%)
 rename airflow/www/static/js/tree/{ => utils}/useErrorToast.js (100%)
 rename airflow/www/static/js/tree/{context/selection.jsx => utils/useSelection.jsx} (50%)

diff --git a/airflow/api_connexion/schemas/task_schema.py b/airflow/api_connexion/schemas/task_schema.py
index 600d28ac0d6c6..aa9a4703088b0 100644
--- a/airflow/api_connexion/schemas/task_schema.py
+++ b/airflow/api_connexion/schemas/task_schema.py
@@ -58,6 +58,7 @@ class TaskSchema(Schema):
     sub_dag = fields.Nested(DAGSchema, dump_only=True)
     downstream_task_ids = fields.List(fields.String(), dump_only=True)
     params = fields.Method('get_params', dump_only=True)
+    is_mapped = fields.Boolean(dump_only=True)
 
     def _get_class_reference(self, obj):
         result = ClassReferenceSchema().dump(obj)
diff --git a/airflow/www/static/js/tree/Tree.jsx b/airflow/www/static/js/tree/Tree.jsx
index a8f39f6a3016c..1283d14c1b938 100644
--- a/airflow/www/static/js/tree/Tree.jsx
+++ b/airflow/www/static/js/tree/Tree.jsx
@@ -39,7 +39,7 @@ import renderTaskRows from './renderTaskRows';
 import ResetRoot from './ResetRoot';
 import DagRuns from './dagRuns';
 import Details from './details';
-import { useSelection } from './context/selection';
+import useSelection from './utils/useSelection';
 import { useAutoRefresh } from './context/autorefresh';
 
 const sidePanelKey = 'hideSidePanel';
diff --git a/airflow/www/static/js/tree/api/useClearRun.js b/airflow/www/static/js/tree/api/useClearRun.js
index afcb620c3aff3..52dcb21e5cfad 100644
--- a/airflow/www/static/js/tree/api/useClearRun.js
+++ b/airflow/www/static/js/tree/api/useClearRun.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const csrfToken = getMetaValue('csrf_token');
 const clearRunUrl = getMetaValue('dagrun_clear_url');
diff --git a/airflow/www/static/js/tree/api/useClearTask.js b/airflow/www/static/js/tree/api/useClearTask.js
index 777a621aaf25c..0733b402b4bbe 100644
--- a/airflow/www/static/js/tree/api/useClearTask.js
+++ b/airflow/www/static/js/tree/api/useClearTask.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const csrfToken = getMetaValue('csrf_token');
 const clearUrl = getMetaValue('clear_url');
diff --git a/airflow/www/static/js/tree/api/useConfirmMarkTask.js b/airflow/www/static/js/tree/api/useConfirmMarkTask.js
index 7db31ab8a879f..a095f5a674a71 100644
--- a/airflow/www/static/js/tree/api/useConfirmMarkTask.js
+++ b/airflow/www/static/js/tree/api/useConfirmMarkTask.js
@@ -20,7 +20,7 @@
 import axios from 'axios';
 import { useMutation } from 'react-query';
 import { getMetaValue } from '../../utils';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const confirmUrl = getMetaValue('confirm_url');
 
diff --git a/airflow/www/static/js/tree/api/useMarkFailedRun.js b/airflow/www/static/js/tree/api/useMarkFailedRun.js
index c7487be5e901d..697a49aa4dc34 100644
--- a/airflow/www/static/js/tree/api/useMarkFailedRun.js
+++ b/airflow/www/static/js/tree/api/useMarkFailedRun.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const csrfToken = getMetaValue('csrf_token');
 const markFailedUrl = getMetaValue('dagrun_failed_url');
diff --git a/airflow/www/static/js/tree/api/useMarkFailedTask.js b/airflow/www/static/js/tree/api/useMarkFailedTask.js
index 4ef72df1ba9a6..3c8f14c7d0e1b 100644
--- a/airflow/www/static/js/tree/api/useMarkFailedTask.js
+++ b/airflow/www/static/js/tree/api/useMarkFailedTask.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const failedUrl = getMetaValue('failed_url');
 const csrfToken = getMetaValue('csrf_token');
diff --git a/airflow/www/static/js/tree/api/useMarkSuccessRun.js b/airflow/www/static/js/tree/api/useMarkSuccessRun.js
index 6076e6ba1e124..0d171e2ed8412 100644
--- a/airflow/www/static/js/tree/api/useMarkSuccessRun.js
+++ b/airflow/www/static/js/tree/api/useMarkSuccessRun.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const markSuccessUrl = getMetaValue('dagrun_success_url');
 const csrfToken = getMetaValue('csrf_token');
diff --git a/airflow/www/static/js/tree/api/useMarkSuccessTask.js b/airflow/www/static/js/tree/api/useMarkSuccessTask.js
index 14d83e653d110..41d68fe24cc48 100644
--- a/airflow/www/static/js/tree/api/useMarkSuccessTask.js
+++ b/airflow/www/static/js/tree/api/useMarkSuccessTask.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const csrfToken = getMetaValue('csrf_token');
 const successUrl = getMetaValue('success_url');
diff --git a/airflow/www/static/js/tree/api/useQueueRun.js b/airflow/www/static/js/tree/api/useQueueRun.js
index 5848d680f925b..0acae341107a7 100644
--- a/airflow/www/static/js/tree/api/useQueueRun.js
+++ b/airflow/www/static/js/tree/api/useQueueRun.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const csrfToken = getMetaValue('csrf_token');
 const queuedUrl = getMetaValue('dagrun_queued_url');
diff --git a/airflow/www/static/js/tree/api/useRunTask.js b/airflow/www/static/js/tree/api/useRunTask.js
index 810b2dceab930..4616d2b1232c8 100644
--- a/airflow/www/static/js/tree/api/useRunTask.js
+++ b/airflow/www/static/js/tree/api/useRunTask.js
@@ -21,7 +21,7 @@ import axios from 'axios';
 import { useMutation, useQueryClient } from 'react-query';
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import useErrorToast from '../useErrorToast';
+import useErrorToast from '../utils/useErrorToast';
 
 const csrfToken = getMetaValue('csrf_token');
 const runUrl = getMetaValue('run_url');
diff --git a/airflow/www/static/js/tree/api/useTreeData.js b/airflow/www/static/js/tree/api/useTreeData.js
index 84886249c450a..9cb9737fd3143 100644
--- a/airflow/www/static/js/tree/api/useTreeData.js
+++ b/airflow/www/static/js/tree/api/useTreeData.js
@@ -24,8 +24,8 @@ import axios from 'axios';
 
 import { getMetaValue } from '../../utils';
 import { useAutoRefresh } from '../context/autorefresh';
-import { formatData, areActiveRuns } from '../treeDataUtils';
-import useErrorToast from '../useErrorToast';
+import { formatData, areActiveRuns } from '../utils/treeData';
+import useErrorToast from '../utils/useErrorToast';
 
 // dagId comes from dag.html
 const dagId = getMetaValue('dag_id');
diff --git a/airflow/www/static/js/tree/context/autorefresh.jsx b/airflow/www/static/js/tree/context/autorefresh.jsx
index 77ca7f342f57e..dc21552028607 100644
--- a/airflow/www/static/js/tree/context/autorefresh.jsx
+++ b/airflow/www/static/js/tree/context/autorefresh.jsx
@@ -21,7 +21,7 @@
 
 import React, { useContext, useState, useEffect } from 'react';
 import { getMetaValue } from '../../utils';
-import { formatData, areActiveRuns } from '../treeDataUtils';
+import { formatData, areActiveRuns } from '../utils/treeData';
 
 const autoRefreshKey = 'disabledAutoRefresh';
 
diff --git a/airflow/www/static/js/tree/dagRuns/index.jsx b/airflow/www/static/js/tree/dagRuns/index.jsx
index 4632cede89843..c304b18bcf402 100644
--- a/airflow/www/static/js/tree/dagRuns/index.jsx
+++ b/airflow/www/static/js/tree/dagRuns/index.jsx
@@ -29,7 +29,7 @@ import {
 import { useTreeData } from '../api';
 import DagRunBar from './Bar';
 import { getDuration, formatDuration } from '../../datetime_utils';
-import { useSelection } from '../context/selection';
+import useSelection from '../utils/useSelection';
 
 const DurationTick = ({ children, ...rest }) => (
   <Text fontSize="sm" color="gray.400" right={1} position="absolute" whiteSpace="nowrap" {...rest}>
diff --git a/airflow/www/static/js/tree/dagRuns/index.test.jsx b/airflow/www/static/js/tree/dagRuns/index.test.jsx
index 5faa8b6e9526b..15a30ce41a5dc 100644
--- a/airflow/www/static/js/tree/dagRuns/index.test.jsx
+++ b/airflow/www/static/js/tree/dagRuns/index.test.jsx
@@ -24,10 +24,10 @@ import { render } from '@testing-library/react';
 import { ChakraProvider, Table, Tbody } from '@chakra-ui/react';
 import moment from 'moment-timezone';
 import { QueryClient, QueryClientProvider } from 'react-query';
+import { MemoryRouter } from 'react-router-dom';
 
 import DagRuns from './index';
 import { ContainerRefProvider } from '../context/containerRef';
-import { SelectionProvider } from '../context/selection';
 import { TimezoneProvider } from '../context/timezone';
 import { AutoRefreshProvider } from '../context/autorefresh';
 
@@ -42,13 +42,13 @@ const Wrapper = ({ children }) => {
           <ContainerRefProvider value={{}}>
             <TimezoneProvider value={{ timezone: 'UTC' }}>
               <AutoRefreshProvider value={{ isRefreshOn: false, stopRefresh: () => {} }}>
-                <SelectionProvider value={{ onSelect: () => {}, selected: {} }}>
+                <MemoryRouter>
                   <Table>
                     <Tbody>
                       {children}
                     </Tbody>
                   </Table>
-                </SelectionProvider>
+                </MemoryRouter>
               </AutoRefreshProvider>
             </TimezoneProvider>
           </ContainerRefProvider>
diff --git a/airflow/www/static/js/tree/details/Header.jsx b/airflow/www/static/js/tree/details/Header.jsx
index 14187f199eaf8..5ce3b0583d71f 100644
--- a/airflow/www/static/js/tree/details/Header.jsx
+++ b/airflow/www/static/js/tree/details/Header.jsx
@@ -29,9 +29,9 @@ import {
 import { MdPlayArrow } from 'react-icons/md';
 
 import { getMetaValue } from '../../utils';
-import { useSelection } from '../context/selection';
+import useSelection from '../utils/useSelection';
 import Time from '../Time';
-import { useTreeData } from '../api';
+import { useTasks, useTreeData } from '../api';
 
 const dagId = getMetaValue('dag_id');
 
@@ -44,8 +44,10 @@ const LabelValue = ({ label, value }) => (
 
 const Header = () => {
   const { data: { dagRuns = [] } } = useTreeData();
-  const { selected: { taskId, runId, task }, onSelect, clearSelection } = useSelection();
+  const { selected: { taskId, runId }, onSelect, clearSelection } = useSelection();
+  const { data: { tasks } } = useTasks();
   const dagRun = dagRuns.find((r) => r.runId === runId);
+  const task = tasks.find((t) => t.taskId === taskId);
 
   let runLabel;
   if (dagRun) {
diff --git a/airflow/www/static/js/tree/details/index.jsx b/airflow/www/static/js/tree/details/index.jsx
index ffe8b0ff74179..ee0fef2dd99ce 100644
--- a/airflow/www/static/js/tree/details/index.jsx
+++ b/airflow/www/static/js/tree/details/index.jsx
@@ -28,7 +28,7 @@ import Header from './Header';
 import TaskInstanceContent from './content/taskInstance';
 import DagRunContent from './content/dagRun';
 import DagContent from './content/Dag';
-import { useSelection } from '../context/selection';
+import useSelection from '../utils/useSelection';
 
 const Details = () => {
   const { selected } = useSelection();
diff --git a/airflow/www/static/js/tree/index.jsx b/airflow/www/static/js/tree/index.jsx
index 5deca3a1c6d40..abe135bebbab7 100644
--- a/airflow/www/static/js/tree/index.jsx
+++ b/airflow/www/static/js/tree/index.jsx
@@ -28,7 +28,6 @@ import createCache from '@emotion/cache';
 import { QueryClient, QueryClientProvider } from 'react-query';
 
 import Tree from './Tree';
-import { SelectionProvider } from './context/selection';
 import { ContainerRefProvider } from './context/containerRef';
 import { TimezoneProvider } from './context/timezone';
 import { AutoRefreshProvider } from './context/autorefresh';
@@ -79,9 +78,7 @@ function App() {
               <TimezoneProvider>
                 <AutoRefreshProvider>
                   <BrowserRouter>
-                    <SelectionProvider>
-                      <Tree />
-                    </SelectionProvider>
+                    <Tree />
                   </BrowserRouter>
                 </AutoRefreshProvider>
               </TimezoneProvider>
diff --git a/airflow/www/static/js/tree/renderTaskRows.jsx b/airflow/www/static/js/tree/renderTaskRows.jsx
index 36fc64c19584f..8713c3801d112 100644
--- a/airflow/www/static/js/tree/renderTaskRows.jsx
+++ b/airflow/www/static/js/tree/renderTaskRows.jsx
@@ -34,7 +34,7 @@ import StatusBox, { boxSize, boxSizePx } from './StatusBox';
 import TaskName from './TaskName';
 
 import { getMetaValue } from '../utils';
-import { useSelection } from './context/selection';
+import useSelection from './utils/useSelection';
 
 const boxPadding = 3;
 const boxPaddingPx = `${boxPadding}px`;
diff --git a/airflow/www/static/js/tree/renderTaskRows.test.jsx b/airflow/www/static/js/tree/renderTaskRows.test.jsx
index b163d62363bfc..88afb4a580c81 100644
--- a/airflow/www/static/js/tree/renderTaskRows.test.jsx
+++ b/airflow/www/static/js/tree/renderTaskRows.test.jsx
@@ -24,10 +24,10 @@ import { render, fireEvent } from '@testing-library/react';
 import { ChakraProvider, Table, Tbody } from '@chakra-ui/react';
 import moment from 'moment';
 import { QueryClient, QueryClientProvider } from 'react-query';
+import { MemoryRouter } from 'react-router-dom';
 
 import renderTaskRows from './renderTaskRows';
 import { ContainerRefProvider } from './context/containerRef';
-import { SelectionProvider } from './context/selection';
 
 global.moment = moment;
 
@@ -101,13 +101,13 @@ const Wrapper = ({ children }) => {
       <ChakraProvider>
         <QueryClientProvider client={queryClient}>
           <ContainerRefProvider value={{}}>
-            <SelectionProvider value={{ onSelect: () => {}, selected: {} }}>
+            <MemoryRouter>
               <Table>
                 <Tbody>
                   {children}
                 </Tbody>
               </Table>
-            </SelectionProvider>
+            </MemoryRouter>
           </ContainerRefProvider>
         </QueryClientProvider>
       </ChakraProvider>
diff --git a/airflow/www/static/js/tree/treeDataUtils.js b/airflow/www/static/js/tree/utils/treeData.js
similarity index 100%
rename from airflow/www/static/js/tree/treeDataUtils.js
rename to airflow/www/static/js/tree/utils/treeData.js
diff --git a/airflow/www/static/js/tree/useErrorToast.js b/airflow/www/static/js/tree/utils/useErrorToast.js
similarity index 100%
rename from airflow/www/static/js/tree/useErrorToast.js
rename to airflow/www/static/js/tree/utils/useErrorToast.js
diff --git a/airflow/www/static/js/tree/context/selection.jsx b/airflow/www/static/js/tree/utils/useSelection.jsx
similarity index 50%
rename from airflow/www/static/js/tree/context/selection.jsx
rename to airflow/www/static/js/tree/utils/useSelection.jsx
index 148ded29585f4..e900d0a6722a6 100644
--- a/airflow/www/static/js/tree/context/selection.jsx
+++ b/airflow/www/static/js/tree/utils/useSelection.jsx
@@ -17,46 +17,19 @@
  * under the License.
  */
 
-import React, {
-  useContext, useReducer, useEffect, useRef,
-} from 'react';
 import { useSearchParams } from 'react-router-dom';
 
-const SelectionContext = React.createContext(null);
-
-const SELECT = 'SELECT';
-const DESELECT = 'DESELECT';
 const RUN_ID = 'dag_run_id';
 const TASK_ID = 'task_id';
 
-const selectionReducer = (state, { type, payload }) => {
-  switch (type) {
-    case SELECT:
-      // Deselect if it is the same selection
-      if (payload.taskId === state.taskId && payload.runId === state.runId) {
-        return {};
-      }
-      return payload;
-    case DESELECT:
-      return {};
-    default:
-      return state;
-  }
-};
-
-// Expose the grid selection to any react component instead of passing around lots of props
-export const SelectionProvider = ({ children }) => {
+const useSelection = () => {
   const [searchParams, setSearchParams] = useSearchParams();
-  const [selected, dispatch] = useReducer(selectionReducer, {});
-
-  const rendered = useRef(false);
 
   // Clear selection, but keep other search params
   const clearSelection = () => {
     searchParams.delete(RUN_ID);
     searchParams.delete(TASK_ID);
     setSearchParams(searchParams);
-    dispatch({ type: DESELECT });
   };
 
   const onSelect = (payload) => {
@@ -64,26 +37,13 @@ export const SelectionProvider = ({ children }) => {
     if (payload.runId) params.set(RUN_ID, payload.runId);
     if (payload.taskId) params.set(TASK_ID, payload.taskId);
     setSearchParams(params);
-    dispatch({ type: SELECT, payload });
   };
 
-  // Check search params and set selection but only on the first render
-  useEffect(() => {
-    if (!rendered.current) {
-      const runId = searchParams.get(RUN_ID);
-      const taskId = searchParams.get(TASK_ID);
-      if (runId) {
-        dispatch({ type: SELECT, payload: { runId, taskId } });
-      }
-      rendered.current = true;
-    }
-  }, [searchParams]);
+  const runId = searchParams.get(RUN_ID);
+  const taskId = searchParams.get(TASK_ID);
+  const selected = { runId, taskId };
 
-  return (
-    <SelectionContext.Provider value={{ selected, clearSelection, onSelect }}>
-      {children}
-    </SelectionContext.Provider>
-  );
+  return { selected, clearSelection, onSelect };
 };
 
-export const useSelection = () => useContext(SelectionContext);
+export default useSelection;

From 032cfd645e3472b4e64b4de3494a7b0d422ec3fe Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Wed, 27 Apr 2022 13:08:00 -0400
Subject: [PATCH 4/5] delete extra params

---
 airflow/www/static/js/tree/utils/useSelection.jsx | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/airflow/www/static/js/tree/utils/useSelection.jsx b/airflow/www/static/js/tree/utils/useSelection.jsx
index e900d0a6722a6..c7220b4f215a5 100644
--- a/airflow/www/static/js/tree/utils/useSelection.jsx
+++ b/airflow/www/static/js/tree/utils/useSelection.jsx
@@ -34,8 +34,13 @@ const useSelection = () => {
 
   const onSelect = (payload) => {
     const params = new URLSearchParams(searchParams);
+
     if (payload.runId) params.set(RUN_ID, payload.runId);
+    else params.delete(RUN_ID);
+
     if (payload.taskId) params.set(TASK_ID, payload.taskId);
+    else params.delete(TASK_ID);
+
     setSearchParams(params);
   };
 

From dbb8c7d560039d480083b110477a49511ed80e0b Mon Sep 17 00:00:00 2001
From: Brent Bovenzi <brent.bovenzi@gmail.com>
Date: Wed, 27 Apr 2022 13:16:31 -0400
Subject: [PATCH 5/5] remove API change

---
 airflow/api_connexion/schemas/task_schema.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/airflow/api_connexion/schemas/task_schema.py b/airflow/api_connexion/schemas/task_schema.py
index aa9a4703088b0..600d28ac0d6c6 100644
--- a/airflow/api_connexion/schemas/task_schema.py
+++ b/airflow/api_connexion/schemas/task_schema.py
@@ -58,7 +58,6 @@ class TaskSchema(Schema):
     sub_dag = fields.Nested(DAGSchema, dump_only=True)
     downstream_task_ids = fields.List(fields.String(), dump_only=True)
     params = fields.Method('get_params', dump_only=True)
-    is_mapped = fields.Boolean(dump_only=True)
 
     def _get_class_reference(self, obj):
         result = ClassReferenceSchema().dump(obj)
