From bb1d9c55e46e5c64c652183520698935da9f8708 Mon Sep 17 00:00:00 2001
From: Davis Plumlee <davis.plumlee@elastic.co>
Date: Mon, 13 Sep 2021 14:32:27 -0400
Subject: [PATCH] addresses comments

---
 .../timeline_actions/alert_context_menu.tsx   | 26 ++++++------
 .../components/take_action_dropdown/index.tsx | 42 +++++++++++++++++--
 2 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/x-pack/plugins/security_solution/public/detections/components/alerts_table/timeline_actions/alert_context_menu.tsx b/x-pack/plugins/security_solution/public/detections/components/alerts_table/timeline_actions/alert_context_menu.tsx
index adadf970cf5fe..3a815468932f0 100644
--- a/x-pack/plugins/security_solution/public/detections/components/alerts_table/timeline_actions/alert_context_menu.tsx
+++ b/x-pack/plugins/security_solution/public/detections/components/alerts_table/timeline_actions/alert_context_menu.tsx
@@ -105,6 +105,18 @@ const AlertContextMenuComponent: React.FC<AlertContextMenuProps & PropsFromRedux
     );
   }, [disabled, onButtonClick, ariaLabel]);
 
+  const refetchQuery = (newQueries: inputsModel.GlobalQuery[]) => {
+    newQueries.forEach((q) => q.refetch && (q.refetch as inputsModel.Refetch)());
+  };
+
+  const refetchAll = useCallback(() => {
+    if (timelineId === TimelineId.active) {
+      refetchQuery([timelineQuery]);
+    } else {
+      refetchQuery(globalQuery);
+    }
+  }, [timelineId, globalQuery, timelineQuery]);
+
   const {
     exceptionModalType,
     onAddExceptionCancel,
@@ -113,7 +125,7 @@ const AlertContextMenuComponent: React.FC<AlertContextMenuProps & PropsFromRedux
     ruleIndices,
   } = useExceptionModal({
     ruleIndex: ecsRowData?.signal?.rule?.index,
-    refetch,
+    refetch: refetchAll,
     timelineId,
   });
 
@@ -123,18 +135,6 @@ const AlertContextMenuComponent: React.FC<AlertContextMenuProps & PropsFromRedux
     onAddEventFilterClick,
   } = useEventFilterModal();
 
-  const refetchQuery = (newQueries: inputsModel.GlobalQuery[]) => {
-    newQueries.forEach((q) => q.refetch && (q.refetch as inputsModel.Refetch)());
-  };
-
-  const refetchAll = useCallback(() => {
-    if (timelineId === TimelineId.active) {
-      refetchQuery([timelineQuery]);
-    } else {
-      refetchQuery(globalQuery);
-    }
-  }, [timelineId, globalQuery, timelineQuery]);
-
   const { actionItems: statusActionItems } = useAlertsActions({
     alertStatus,
     eventId: ecsRowData?._id,
diff --git a/x-pack/plugins/security_solution/public/detections/components/take_action_dropdown/index.tsx b/x-pack/plugins/security_solution/public/detections/components/take_action_dropdown/index.tsx
index 0432e7d353086..425d049388764 100644
--- a/x-pack/plugins/security_solution/public/detections/components/take_action_dropdown/index.tsx
+++ b/x-pack/plugins/security_solution/public/detections/components/take_action_dropdown/index.tsx
@@ -9,7 +9,8 @@ import React, { useState, useCallback, useMemo } from 'react';
 import { EuiContextMenuPanel, EuiButton, EuiPopover } from '@elastic/eui';
 import type { ExceptionListType } from '@kbn/securitysolution-io-ts-list-types';
 import { isEmpty } from 'lodash/fp';
-import { TimelineEventsDetailsItem } from '../../../../common';
+import { connect, ConnectedProps } from 'react-redux';
+import { TimelineEventsDetailsItem, TimelineId } from '../../../../common';
 import { TAKE_ACTION } from '../alerts_table/alerts_utility_bar/translations';
 import { useExceptionActions } from '../alerts_table/timeline_actions/use_add_exception_actions';
 import { useAlertsActions } from '../alerts_table/timeline_actions/use_alerts_actions';
@@ -23,6 +24,7 @@ import { Status } from '../../../../common/detection_engine/schemas/common/schem
 import { isAlertFromEndpointAlert } from '../../../common/utils/endpoint_alert_check';
 import { useIsExperimentalFeatureEnabled } from '../../../common/hooks/use_experimental_features';
 import { useAddToCaseActions } from '../alerts_table/timeline_actions/use_add_to_case_actions';
+import { inputsModel, inputsSelectors, State } from '../../../common/store';
 
 interface ActionsData {
   alertStatus: Status;
@@ -46,7 +48,7 @@ export interface TakeActionDropdownProps {
   timelineId: string;
 }
 
-export const TakeActionDropdown = React.memo(
+export const TakeActionDropdownComponent = React.memo(
   ({
     detailsData,
     ecsData,
@@ -59,7 +61,9 @@ export const TakeActionDropdown = React.memo(
     onAddIsolationStatusClick,
     refetch,
     timelineId,
-  }: TakeActionDropdownProps) => {
+    globalQuery,
+    timelineQuery,
+  }: TakeActionDropdownProps & PropsFromRedux) => {
     const tGridEnabled = useIsExperimentalFeatureEnabled('tGridEnabled');
 
     const [isPopoverOpen, setIsPopoverOpen] = useState(false);
@@ -141,12 +145,24 @@ export const TakeActionDropdown = React.memo(
       closePopoverHandler();
     }, [closePopoverHandler]);
 
+    const refetchQuery = (newQueries: inputsModel.GlobalQuery[]) => {
+      newQueries.forEach((q) => q.refetch && (q.refetch as inputsModel.Refetch)());
+    };
+
+    const refetchAll = useCallback(() => {
+      if (timelineId === TimelineId.active) {
+        refetchQuery([timelineQuery]);
+      } else {
+        refetchQuery(globalQuery);
+      }
+    }, [timelineId, globalQuery, timelineQuery]);
+
     const { actionItems: statusActionItems } = useAlertsActions({
       alertStatus: actionsData.alertStatus,
       closePopover: closePopoverAndFlyout,
       eventId: actionsData.eventId,
       indexName,
-      refetch,
+      refetch: refetchAll,
       timelineId,
     });
 
@@ -216,3 +232,21 @@ export const TakeActionDropdown = React.memo(
     ) : null;
   }
 );
+
+const makeMapStateToProps = () => {
+  const getGlobalQueries = inputsSelectors.globalQuery();
+  const getTimelineQuery = inputsSelectors.timelineQueryByIdSelector();
+  const mapStateToProps = (state: State, { timelineId }: TakeActionDropdownProps) => {
+    return {
+      globalQuery: getGlobalQueries(state),
+      timelineQuery: getTimelineQuery(state, timelineId),
+    };
+  };
+  return mapStateToProps;
+};
+
+const connector = connect(makeMapStateToProps);
+
+type PropsFromRedux = ConnectedProps<typeof connector>;
+
+export const TakeActionDropdown = connector(React.memo(TakeActionDropdownComponent));