Remember expanded task groups in localStorage (#14661)

* Save expanded and focused state of TaskGroups in localStorage

* Restore focus on refresh

* Address style issues according to comments

(cherry picked from commit 456a7dd)

airflow/www/templates/airflow/graph.html

@@ -235,7 +235,7 @@
         });
 
         d3.selectAll("g.node").on("mouseout", function (d) {
-          d3.select(this).selectAll("rect").style("stroke", null);
+          d3.select(this).selectAll("rect,circle").style("stroke", null);
           highlight_nodes(g.predecessors(d), null, initialStrokeWidth)
           highlight_nodes(g.successors(d), null, initialStrokeWidth)
           d3.selectAll("g.node")
@@ -244,6 +244,7 @@
             .style("stroke-width", initialStrokeWidth);
           d3.selectAll("g.edgePath")
             .style("opacity", 1);
+          localStorage.removeItem(focused_group_key(dag_id));
         });
         updateNodesStates(task_instances);
         setUpZoomSupport();
@@ -417,6 +418,8 @@
               .style("opacity", 1);
           d3.selectAll('.js-state-legend-item')
               .style("background-color", null);
+
+          localStorage.removeItem(focused_group_key(dag_id));
       }
 
       function focusState(state, node, color){
@@ -591,6 +594,22 @@
         return children
       }
 
+      // Return list of all task group ids in the given task group including the given group itself.
+      function get_all_group_ids(group) {
+        var children = [group.id];
+
+        for (const [key, val] of Object.entries(group.children)) {
+          if (val.children != undefined) {
+            // group
+            const sub_group_children = get_all_group_ids(val)
+            for (const id of sub_group_children) {
+              children.push(id);
+            }
+          }
+        }
+        return children;
+      }
+
 
       // Return the state for the node based on the state of its taskinstance or that of its
       // children if it's a group node
@@ -626,6 +645,16 @@
         return "no_status"
       }
 
+      // Returns the key used to store expanded task group ids in localStorage
+      function expanded_groups_key(dag_id) {
+          return `expanded_groups_${dag_id}`;
+      }
+
+      // Returns the key used to store the focused task group id in localStorage
+      function focused_group_key(dag_id) {
+          return `focused_group_${dag_id}`;
+      }
+
       // Focus the graph on the expanded/collapsed node
       function focus_group(node_id) {
         if(node_id != null && zoom != null) {
@@ -668,11 +697,13 @@
                       .style("opacity", 0.2).duration(duration)
               }
             });
+
+            localStorage.setItem(focused_group_key(dag_id), node_id);
         }
       }
 
       // Expands a group node
-      function expand_group(node_id, node) {
+      function expand_group(node_id, node, focus=true) {
         node.children.forEach(function (val) {
           // Set children nodes
           g.setNode(val.id, val.value)
@@ -706,17 +737,22 @@
         })
 
         draw()
-        focus_group(node_id)
+
+        if (focus) {
+          focus_group(node_id);
+        }
+
+        save_expanded_group(node_id)
     }
 
     // Remove the node with this node_id from g.
     function remove_node(node_id) {
-      if(g.hasNode(node_id)) {
+      if (g.hasNode(node_id)) {
           node = g.node(node_id)
           if(node.children != undefined) {
             // If the child is an expanded group node, remove children too.
             node.children.forEach(function (child) {
-              remove_node(child.id)
+              remove_node(child.id);
             })
           }
       }
@@ -745,10 +781,77 @@
 
         draw()
         focus_group(node_id)
+
+        remove_expanded_group(node_id, node);
       }
 
-      expand_group(null, nodes)
+    function get_saved_groups(dag_id) {
+        // expanded_groups is a Set
+        try {
+            var expanded_groups = new Set(JSON.parse(localStorage.getItem(expanded_groups_key(dag_id))));
+        } catch {
+            var expanded_groups = new Set();
+        }
+
+        return expanded_groups;
+    }
+
+    // Clean up invalid group_ids from saved_group_ids (e.g. due to DAG changes)
+    function prune_invalid_saved_group_ids() {
+        // All the group_ids in the whole DAG
+        const all_group_ids = new Set(get_all_group_ids(nodes));
+        var expanded_groups = get_saved_groups(dag_id);
+        expanded_groups = Array.from(expanded_groups).filter(group_id => all_group_ids.has(group_id));
+        localStorage.setItem(expanded_groups_key(dag_id), JSON.stringify(expanded_groups));
+    }
+
+    // Remember the expanded groups in local storage so that it can be used to restore the expanded state
+    // of task groups.
+    function save_expanded_group(node_id) {
+        // expanded_groups is a Set
+        var expanded_groups = get_saved_groups(dag_id);
+        expanded_groups.add(node_id)
+        localStorage.setItem(expanded_groups_key(dag_id), JSON.stringify(Array.from(expanded_groups)));
+    }
+
+    // Remove the node_id from the expanded state
+    function remove_expanded_group(node_id, node) {
+        var expanded_groups = get_saved_groups(dag_id);
+        const child_group_ids = get_all_group_ids(node);
+        child_group_ids.forEach(child_id => expanded_groups.delete(child_id));
+        localStorage.setItem(expanded_groups_key(dag_id), JSON.stringify(Array.from(expanded_groups)));
+    }
+
+    // Restore previously expanded task groups
+    function expand_saved_groups(expanded_groups, node) {
+        if (node.children == undefined) {
+            return;
+        }
+
+        node.children.forEach(function (child_node) {
+            if(expanded_groups.has(child_node.id)) {
+                expand_group(child_node.id, g.node(child_node.id), false);
+
+                expand_saved_groups(expanded_groups, child_node);
+            }
+        });
+    }
+
+    prune_invalid_saved_group_ids();
+    const focus_node_id = localStorage.getItem(focused_group_key(dag_id));
+    const expanded_groups = get_saved_groups(dag_id);
+
+    // Always expand the root node
+    expand_group(null, nodes);
+
+    // Expand the node that were previously expanded
+    expand_saved_groups(expanded_groups, nodes);
+
+    // Restore focus (if available)
+    if(g.hasNode(focus_node_id)) {
+      focus_group(focus_node_id);
+    }
 
-      initRefresh();
+    initRefresh();
   </script>
 {% endblock %}