Fleet UI: Can run a live query on an edited (but not saved) existing query (#16282)

changes/16067-bug-running-edited-query

@@ -0,0 +1 @@
+- Ability to run a live query on an edited existing query before saving

frontend/__mocks__/queryMock.ts

@@ -27,6 +27,7 @@ const DEFAULT_QUERY_MOCK: ISchedulableQuery = {
     system_time_p95: 1,
     total_executions: 6,
   },
+  editingExistingQuery: false,
 };
 
 const createMockQuery = (

frontend/__mocks__/scheduleableQueryMock.ts

@@ -29,6 +29,7 @@ const DEFAULT_SCHEDULABLE_QUERY_MOCK: ISchedulableQuery = {
     user_time_p95: 251.4615,
     total_executions: 5746,
   },
+  editingExistingQuery: false,
 };
 
 const createMockSchedulableQuery = (

frontend/context/query.tsx

@@ -29,6 +29,7 @@ type InitialStateType = {
   lastEditedQueryMinOsqueryVersion: string;
   lastEditedQueryLoggingType: QueryLoggingOption;
   lastEditedQueryDiscardData: boolean;
+  editingExistingQuery: boolean;
   selectedQueryTargets: ITarget[]; // Mimicks old selectedQueryTargets still used for policies for SelectTargets.tsx and running a live query
   selectedQueryTargetsByType: ISelectedTargetsByType; // New format by type for cleaner app wide state
   setLastEditedQueryId: (value: number | null) => void;
@@ -44,6 +45,7 @@ type InitialStateType = {
   setSelectedOsqueryTable: (tableName: string) => void;
   setSelectedQueryTargets: (value: ITarget[]) => void;
   setSelectedQueryTargetsByType: (value: ISelectedTargetsByType) => void;
+  setEditingExistingQuery: (value: boolean) => void;
 };
 
 export type IQueryContext = InitialStateType;
@@ -61,6 +63,7 @@ const initialState = {
   lastEditedQueryMinOsqueryVersion: DEFAULT_QUERY.min_osquery_version,
   lastEditedQueryLoggingType: DEFAULT_QUERY.logging,
   lastEditedQueryDiscardData: DEFAULT_QUERY.discard_data,
+  editingExistingQuery: DEFAULT_QUERY.editingExistingQuery,
   selectedQueryTargets: DEFAULT_TARGETS,
   selectedQueryTargetsByType: DEFAULT_TARGETS_BY_TYPE,
   setLastEditedQueryId: () => null,
@@ -76,6 +79,7 @@ const initialState = {
   setSelectedOsqueryTable: () => null,
   setSelectedQueryTargets: () => null,
   setSelectedQueryTargetsByType: () => null,
+  setEditingExistingQuery: () => null,
 };
 
 const actions = {
@@ -137,6 +141,10 @@ const reducer = (state: InitialStateType, action: any) => {
           typeof action.lastEditedQueryDiscardData === "undefined"
             ? state.lastEditedQueryDiscardData
             : action.lastEditedQueryDiscardData,
+        editingExistingQuery:
+          typeof action.editingExistingQuery === "undefined"
+            ? state.editingExistingQuery
+            : action.editingExistingQuery,
       };
     case actions.SET_SELECTED_QUERY_TARGETS:
       return {
@@ -176,6 +184,7 @@ const QueryProvider = ({ children }: Props) => {
     lastEditedQueryMinOsqueryVersion: state.lastEditedQueryMinOsqueryVersion,
     lastEditedQueryLoggingType: state.lastEditedQueryLoggingType,
     lastEditedQueryDiscardData: state.lastEditedQueryDiscardData,
+    editingExistingQuery: state.editingExistingQuery,
     selectedQueryTargets: state.selectedQueryTargets,
     selectedQueryTargetsByType: state.selectedQueryTargetsByType,
     setLastEditedQueryId: (lastEditedQueryId: number | null) => {
@@ -242,6 +251,12 @@ const QueryProvider = ({ children }: Props) => {
         lastEditedQueryDiscardData,
       });
     },
+    setEditingExistingQuery: (editingExistingQuery: boolean) => {
+      dispatch({
+        type: actions.SET_LAST_EDITED_QUERY_INFO,
+        editingExistingQuery,
+      });
+    },
     setSelectedQueryTargets: (selectedQueryTargets: ITarget[]) => {
       dispatch({
         type: actions.SET_SELECTED_QUERY_TARGETS,

frontend/interfaces/schedulable_query.ts

@@ -24,6 +24,7 @@ export interface ISchedulableQuery {
   discard_data: boolean;
   packs: IPack[];
   stats: ISchedulableQueryStats;
+  editingExistingQuery: boolean;
 }
 
 export interface IEnhancedQuery extends ISchedulableQuery {

frontend/pages/queries/edit/EditQueryPage.tsx

@@ -71,6 +71,7 @@ const EditQueryPage = ({
     config,
   } = useContext(AppContext);
   const {
+    editingExistingQuery,
     selectedOsqueryTable,
     setSelectedOsqueryTable,
     lastEditedQueryName,
@@ -127,7 +128,7 @@ const EditQueryPage = ({
     ["query", queryId],
     () => queryAPI.load(queryId as number),
     {
-      enabled: !!queryId,
+      enabled: !!queryId && !editingExistingQuery,
       refetchOnWindowFocus: false,
       select: (data) => data.query,
       onSuccess: (returnedQuery) => {

frontend/pages/queries/edit/components/EditQueryForm/EditQueryForm.tsx

@@ -150,6 +150,7 @@ const EditQueryForm = ({
     setLastEditedQueryMinOsqueryVersion,
     setLastEditedQueryLoggingType,
     setLastEditedQueryDiscardData,
+    setEditingExistingQuery,
   } = useContext(QueryContext);
 
   const {
@@ -825,6 +826,7 @@ const EditQueryForm = ({
                 className={`${baseClass}__run`}
                 variant="blue-green"
                 onClick={() => {
+                  setEditingExistingQuery(true); // Persists edited query data through live query flow
                   router.push(
                     PATHS.LIVE_QUERY(queryIdForEdit) +
                       TAGGED_TEMPLATES.queryByHostRoute(hostId)

frontend/pages/queries/live/LiveQueryPage/LiveQueryPage.tsx

@@ -44,6 +44,7 @@ const RunQueryPage = ({
   const handlePageError = useErrorHandler();
   const { config } = useContext(AppContext);
   const {
+    editingExistingQuery,
     selectedQueryTargets,
     setSelectedQueryTargets,
     selectedQueryTargetsByType,
@@ -88,7 +89,7 @@ const RunQueryPage = ({
     Error,
     ISchedulableQuery
   >(["query", queryId], () => queryAPI.load(queryId as number), {
-    enabled: !!queryId,
+    enabled: !!queryId && !editingExistingQuery,
     refetchOnWindowFocus: false,
     select: (data) => data.query,
     onSuccess: (returnedQuery) => {

frontend/pages/queries/live/screens/RunQuery.tsx

@@ -47,6 +47,8 @@ const RunQuery = ({
     DEFAULT_CAMPAIGN_STATE
   );
 
+  const isStoredQueryEdited = storedQuery?.query !== lastEditedQueryBody;
+
   const ws = useRef(null);
   const runQueryInterval = useRef<any>(null);
   const globalSocket = useRef<any>(null);
@@ -168,8 +170,6 @@ const RunQuery = ({
     destroyCampaign();
 
     try {
-      const isStoredQueryEdited = storedQuery?.query !== lastEditedQueryBody;
-
       const returnedCampaign = await queryAPI.run({
         query: lastEditedQueryBody,
         queryId: isStoredQueryEdited ? null : queryId, // we treat edited SQL as a new query

frontend/utilities/constants.tsx

@@ -132,6 +132,7 @@ export const DEFAULT_QUERY: ISchedulableQuery = {
   team_id: 0,
   author_email: "",
   stats: {},
+  editingExistingQuery: false,
 };
 
 export const DEFAULT_CAMPAIGN = {