From 749c1270df53ee96b39923986dbf5dbd50e95887 Mon Sep 17 00:00:00 2001
From: RachelElysia <71795832+RachelElysia@users.noreply.github.com>
Date: Mon, 6 May 2024 16:40:19 -0400
Subject: [PATCH] [unreleased bug] Fleet UI: Automated queries policies
 (#18774)

---
 .../ManagePoliciesPage/ManagePoliciesPage.tsx | 36 +++++++++++--------
 .../ManageQueriesPage/ManageQueriesPage.tsx   | 28 +++++++++++----
 2 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx b/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx
index 340a769c98..b1a76ce687 100644
--- a/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx
+++ b/frontend/pages/policies/ManagePoliciesPage/ManagePoliciesPage.tsx
@@ -18,6 +18,7 @@ import {
   ILoadAllPoliciesResponse,
   ILoadTeamPoliciesResponse,
   IPoliciesCountResponse,
+  IPolicy,
 } from "interfaces/policy";
 import { ITeamConfig } from "interfaces/team";
 
@@ -133,8 +134,10 @@ const ManagePolicyPage = ({
   const [showAddPolicyModal, setShowAddPolicyModal] = useState(false);
   const [showDeletePolicyModal, setShowDeletePolicyModal] = useState(false);
   const [showCalendarEventsModal, setShowCalendarEventsModal] = useState(false);
-
-  const [teamPolicies, setTeamPolicies] = useState<IPolicyStats[]>();
+  const [
+    policiesAvailableToAutomate,
+    setPoliciesAvailableToAutomate,
+  ] = useState<IPolicyStats[]>([]);
 
   // Functions to avoid race conditions
   const initialSearchQuery = (() => queryParams.query ?? "")();
@@ -214,6 +217,9 @@ const ManagePolicyPage = ({
       enabled: isRouteOk && !isAnyTeamSelected,
       select: (data) => data.policies,
       staleTime: 5000,
+      onSuccess: (data) => {
+        setPoliciesAvailableToAutomate(data);
+      },
     }
   );
 
@@ -240,13 +246,14 @@ const ManagePolicyPage = ({
   );
 
   const {
+    data: teamPolicies,
     error: teamPoliciesError,
     isFetching: isFetchingTeamPolicies,
     refetch: refetchTeamPolicies,
   } = useQuery<
     ILoadTeamPoliciesResponse,
     Error,
-    ILoadTeamPoliciesResponse,
+    IPolicyStats[],
     ITeamPoliciesQueryKey[]
   >(
     [
@@ -266,8 +273,12 @@ const ManagePolicyPage = ({
     },
     {
       enabled: isRouteOk && isPremiumTier && !!teamIdForApi,
+      select: (data: ILoadTeamPoliciesResponse) => data.policies,
       onSuccess: (data) => {
-        setTeamPolicies(data.policies);
+        const allPoliciesAvailableToAutomate = data.filter(
+          (policy: IPolicy) => policy.team_id === currentTeamId
+        );
+        setPoliciesAvailableToAutomate(allPoliciesAvailableToAutomate);
       },
     }
   );
@@ -328,12 +339,10 @@ const ManagePolicyPage = ({
     }
   );
 
-  const canAddOrDeletePolicy: boolean =
+  const canAddOrDeletePolicy =
     isGlobalAdmin || isGlobalMaintainer || isTeamMaintainer || isTeamAdmin;
-  const canManageAutomations: boolean = isGlobalAdmin || isTeamAdmin;
-  const hasPoliciesToAutomateOrDelete: boolean = teamIdForApi
-    ? !isFetchingTeamCount && !!teamPoliciesCount && teamPoliciesCount > 0
-    : !!globalPoliciesCount && globalPoliciesCount > 0;
+  const canManageAutomations = isGlobalAdmin || isTeamAdmin;
+  const hasPoliciesToAutomateOrDelete = policiesAvailableToAutomate.length > 0;
 
   const {
     data: config,
@@ -513,7 +522,7 @@ const ManagePolicyPage = ({
 
       // update changed policies calendar events enabled
       const changedPolicies = formData.policies.filter((formPolicy) => {
-        const prevPolicyState = teamPolicies?.find(
+        const prevPolicyState = policiesAvailableToAutomate.find(
           (policy) => policy.id === formPolicy.id
         );
         return (
@@ -588,9 +597,6 @@ const ManagePolicyPage = ({
     }
   };
 
-  const availablePoliciesForAutomation =
-    (isAnyTeamSelected ? teamPolicies : globalPolicies) || [];
-
   const policiesErrors = isAnyTeamSelected
     ? teamPoliciesError
     : globalPoliciesError;
@@ -801,7 +807,7 @@ const ManagePolicyPage = ({
           <OtherWorkflowsModal
             automationsConfig={automationsConfig}
             availableIntegrations={config.integrations}
-            availablePolicies={availablePoliciesForAutomation}
+            availablePolicies={policiesAvailableToAutomate}
             isUpdatingAutomations={isUpdatingAutomations}
             onExit={toggleOtherWorkflowsModal}
             handleSubmit={handleUpdateOtherWorkflows}
@@ -834,7 +840,7 @@ const ManagePolicyPage = ({
                 ?.enable_calendar_events ?? false
             }
             url={teamConfig?.integrations.google_calendar?.webhook_url || ""}
-            policies={teamPolicies || []}
+            policies={policiesAvailableToAutomate}
             isUpdating={updatingPolicyEnabledCalendarEvents}
           />
         )}
diff --git a/frontend/pages/queries/ManageQueriesPage/ManageQueriesPage.tsx b/frontend/pages/queries/ManageQueriesPage/ManageQueriesPage.tsx
index a6a4bd7d8b..c56da46073 100644
--- a/frontend/pages/queries/ManageQueriesPage/ManageQueriesPage.tsx
+++ b/frontend/pages/queries/ManageQueriesPage/ManageQueriesPage.tsx
@@ -117,6 +117,9 @@ const ManageQueriesPage = ({
   const [showPreviewDataModal, setShowPreviewDataModal] = useState(false);
   const [isUpdatingQueries, setIsUpdatingQueries] = useState(false);
   const [isUpdatingAutomations, setIsUpdatingAutomations] = useState(false);
+  const [queriesAvailableToAutomate, setQueriesAvailableToAutomate] = useState<
+    IEnhancedQuery[] | []
+  >([]);
 
   const {
     data: enhancedQueries,
@@ -138,6 +141,19 @@ const ManageQueriesPage = ({
       refetchOnWindowFocus: false,
       enabled: isRouteOk,
       staleTime: 5000,
+      onSuccess: (data) => {
+        if (data) {
+          const enhancedAllQueries = data.map(enhanceQuery);
+
+          const allQueriesAvailableToAutomate = teamIdForApi
+            ? enhancedAllQueries.filter(
+                (query: IEnhancedQuery) => query.team_id === currentTeamId
+              )
+            : enhancedAllQueries;
+
+          setQueriesAvailableToAutomate(allQueriesAvailableToAutomate);
+        }
+      },
     }
   );
 
@@ -150,12 +166,10 @@ const ManageQueriesPage = ({
   }, [teamIdForApi, enhancedQueries]);
 
   const automatedQueryIds = useMemo(() => {
-    return enhancedQueries
-      ? enhancedQueries
-          .filter((query) => query.automations_enabled)
-          .map((query) => query.id)
-      : [];
-  }, [enhancedQueries]);
+    return queriesAvailableToAutomate
+      .filter((query) => query.automations_enabled)
+      .map((query) => query.id);
+  }, [queriesAvailableToAutomate]);
 
   useEffect(() => {
     const path = location.pathname + location.search;
@@ -330,7 +344,7 @@ const ManageQueriesPage = ({
             onCancel={toggleManageAutomationsModal}
             isShowingPreviewDataModal={showPreviewDataModal}
             togglePreviewDataModal={togglePreviewDataModal}
-            availableQueries={enhancedQueries}
+            availableQueries={queriesAvailableToAutomate}
             automatedQueryIds={automatedQueryIds}
             logDestination={config?.logging.result.plugin || ""}
           />