datatrails: detect if current trail state is bookmarked (#83283)

* fix: detect if current trail state is bookmarked

public/app/features/trails/DataTrail.tsx

@@ -86,6 +86,8 @@ export class DataTrail extends SceneObjectBase<DataTrailState> {
       const newStepWasAppended = newNumberOfSteps > oldNumberOfSteps;
 
       if (newStepWasAppended) {
+        // In order for the `useBookmarkState` to re-evaluate after a new step was made:
+        this.forceRender();
         // Do nothing because the state is already up to date -- it created a new step!
         return;
       }

public/app/features/trails/MetricScene.tsx

@@ -1,5 +1,5 @@
 import { css } from '@emotion/css';
-import React, { useState } from 'react';
+import React from 'react';
 
 import { DashboardCursorSync, GrafanaTheme2 } from '@grafana/data';
 import {
@@ -26,7 +26,7 @@ import { getAutoQueriesForMetric } from './AutomaticMetricQueries/AutoQueryEngin
 import { AutoVizPanel } from './AutomaticMetricQueries/AutoVizPanel';
 import { AutoQueryDef, AutoQueryInfo } from './AutomaticMetricQueries/types';
 import { ShareTrailButton } from './ShareTrailButton';
-import { getTrailStore } from './TrailStore/TrailStore';
+import { useBookmarkState } from './TrailStore/useBookmarkState';
 import {
   ActionViewDefinition,
   ActionViewType,
@@ -151,14 +151,9 @@ export class MetricActionBar extends SceneObjectBase<MetricActionBarState> {
     const metricScene = sceneGraph.getAncestor(model, MetricScene);
     const styles = useStyles2(getStyles);
     const trail = getTrailFor(model);
-    const [isBookmarked, setBookmarked] = useState(false);
+    const [isBookmarked, toggleBookmark] = useBookmarkState(trail);
     const { actionView } = metricScene.useState();
 
-    const onBookmarkTrail = () => {
-      getTrailStore().addBookmark(trail);
-      setBookmarked(!isBookmarked);
-    };
-
     return (
       <Box paddingY={1}>
         <div className={styles.actions}>
@@ -180,7 +175,7 @@ export class MetricActionBar extends SceneObjectBase<MetricActionBarState> {
                 )
               }
               tooltip={'Bookmark'}
-              onClick={onBookmarkTrail}
+              onClick={toggleBookmark}
             />
             {trail.state.embedded && (
               <ToolbarButton variant={'canvas'} onClick={model.onOpenTrail}>

public/app/features/trails/TrailStore/TrailStore.test.ts

@@ -167,7 +167,7 @@ describe('TrailStore', () => {
     });
   });
   describe('Initialize store with one bookmark trail', () => {
-    beforeAll(() => {
+    beforeEach(() => {
       localStorage.clear();
       localStorage.setItem(
         BOOKMARKED_TRAILS_KEY,
@@ -225,6 +225,35 @@ describe('TrailStore', () => {
       expect(store.recent.length).toBe(1);
     });
 
+    it('should be able to obtain index of bookmark', () => {
+      const trail = store.bookmarks[0].resolve();
+      const index = store.getBookmarkIndex(trail);
+      expect(index).toBe(0);
+    });
+
+    it('index should be undefined for removed bookmarks', () => {
+      const trail = store.bookmarks[0].resolve();
+      store.removeBookmark(0);
+      const index = store.getBookmarkIndex(trail);
+      expect(index).toBe(undefined);
+    });
+
+    it('index should be undefined for a trail that has changed since it was bookmarked', () => {
+      const trail = store.bookmarks[0].resolve();
+      trail.setState({ metric: 'something_completely_different' });
+      const index = store.getBookmarkIndex(trail);
+      expect(index).toBe(undefined);
+    });
+
+    it('should be able to obtain index of a bookmark for a trail that changed back to bookmarked state', () => {
+      const trail = store.bookmarks[0].resolve();
+      const bookmarkedMetric = trail.state.metric;
+      trail.setState({ metric: 'something_completely_different' });
+      trail.setState({ metric: bookmarkedMetric });
+      const index = store.getBookmarkIndex(trail);
+      expect(index).toBe(0);
+    });
+
     it('should remove a bookmark', () => {
       expect(store.bookmarks.length).toBe(1);
       store.removeBookmark(0);

public/app/features/trails/TrailStore/TrailStore.ts

@@ -100,6 +100,7 @@ export class TrailStore {
   load() {
     this._recent = this._loadFromStorage(RECENT_TRAILS_KEY);
     this._bookmarks = this._loadFromStorage(BOOKMARKED_TRAILS_KEY);
+    this._refreshBookmarkIndexMap();
   }
 
   setRecentTrail(trail: DataTrail) {
@@ -125,15 +126,47 @@ export class TrailStore {
 
   addBookmark(trail: DataTrail) {
     this._bookmarks.unshift(trail.getRef());
+    this._refreshBookmarkIndexMap();
     this._save();
   }
 
   removeBookmark(index: number) {
     if (index < this._bookmarks.length) {
       this._bookmarks.splice(index, 1);
+      this._refreshBookmarkIndexMap();
       this._save();
     }
   }
+
+  getBookmarkIndex(trail: DataTrail) {
+    const bookmarkKey = getBookmarkKey(trail);
+    const bookmarkIndex = this._bookmarkIndexMap.get(bookmarkKey);
+    return bookmarkIndex;
+  }
+
+  private _bookmarkIndexMap = new Map<string, number>();
+
+  private _refreshBookmarkIndexMap() {
+    this._bookmarkIndexMap.clear();
+    this._bookmarks.forEach((bookmarked, index) => {
+      const trail = bookmarked.resolve();
+      const key = getBookmarkKey(trail);
+      // If there are duplicate bookmarks, the latest index will be kept
+      this._bookmarkIndexMap.set(key, index);
+    });
+  }
+}
+
+function getBookmarkKey(trail: DataTrail) {
+  const urlState = getUrlSyncManager().getUrlState(trail);
+  // Not part of state
+  delete urlState.actionView;
+  // Populate defaults
+  if (urlState['var-groupby'] === '') {
+    urlState['var-groupby'] = '$__all';
+  }
+  const key = JSON.stringify(urlState);
+  return key;
 }
 
 let store: TrailStore | undefined;

public/app/features/trails/TrailStore/useBookmarkState.ts

@@ -0,0 +1,39 @@
+import { useState } from 'react';
+
+import { DataTrail } from '../DataTrail';
+
+import { getTrailStore } from './TrailStore';
+
+export function useBookmarkState(trail: DataTrail) {
+  // Note that trail object may stay the same, but the state used by `getBookmarkIndex` result may
+  // differ for each re-render of this hook
+  const getBookmarkIndex = () => getTrailStore().getBookmarkIndex(trail);
+
+  const indexOnRender = getBookmarkIndex();
+
+  const [bookmarkIndex, setBookmarkIndex] = useState(indexOnRender);
+
+  // Check if index changed and force a re-render
+  if (indexOnRender !== bookmarkIndex) {
+    setBookmarkIndex(indexOnRender);
+  }
+
+  const isBookmarked = bookmarkIndex != null;
+
+  const toggleBookmark = () => {
+    if (isBookmarked) {
+      let indexToRemove = getBookmarkIndex();
+      while (indexToRemove != null) {
+        // This loop will remove all indices that have an equivalent bookmark key
+        getTrailStore().removeBookmark(indexToRemove);
+        indexToRemove = getBookmarkIndex();
+      }
+    } else {
+      getTrailStore().addBookmark(trail);
+    }
+    setBookmarkIndex(getBookmarkIndex());
+  };
+
+  const result: [typeof isBookmarked, typeof toggleBookmark] = [isBookmarked, toggleBookmark];
+  return result;
+}