Support cancellation of async queries (#1177) (#1186)

* update routes and services



* update public components



* fix imports



---------


(cherry picked from commit 2d229b1)

Signed-off-by: Shenoy Pratik <sgguruda@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

public/components/common/search/sql_search.tsx

@@ -21,11 +21,15 @@ import {
 import { isEqual } from 'lodash';
 import React, { useEffect, useState } from 'react';
 import { batch, useDispatch, useSelector } from 'react-redux';
+import { QUERY_LANGUAGE } from '../../../../common/constants/data_sources';
 import { APP_ANALYTICS_TAB_ID_REGEX, RAW_QUERY } from '../../../../common/constants/explorer';
-import { DirectQueryLoadingStatus, DirectQueryRequest } from '../../../../common/types/explorer';
 import { PPL_NEWLINE_REGEX, PPL_SPAN_REGEX } from '../../../../common/constants/shared';
+import { DirectQueryLoadingStatus, DirectQueryRequest } from '../../../../common/types/explorer';
 import { uiSettingsService } from '../../../../common/utils';
+import { getAsyncSessionId, setAsyncSessionId } from '../../../../common/utils/query_session_utils';
+import { get as getObjValue } from '../../../../common/utils/shared';
 import { useFetchEvents } from '../../../components/event_analytics/hooks';
+import { changeQuery } from '../../../components/event_analytics/redux/slices/query_slice';
 import { usePolling } from '../../../components/hooks/use_polling';
 import { coreRefs } from '../../../framework/core_refs';
 import { SQLService } from '../../../services/requests/sql';
@@ -36,10 +40,6 @@ import {
 } from '../../event_analytics/redux/slices/search_meta_data_slice';
 import { PPLReferenceFlyout } from '../helpers';
 import { Autocomplete } from './autocomplete';
-import { changeQuery } from '../../../components/event_analytics/redux/slices/query_slice';
-import { QUERY_LANGUAGE } from '../../../../common/constants/data_sources';
-import { getAsyncSessionId, setAsyncSessionId } from '../../../../common/utils/query_session_utils';
-import { get as getObjValue } from '../../../../common/utils/shared';
 export interface IQueryBarProps {
   query: string;
   tempQuery: string;
@@ -208,6 +208,7 @@ export const DirectSearch = (props: any) => {
       .then((result) => {
         setAsyncSessionId(getObjValue(result, 'sessionId', null));
         if (result.queryId) {
+          dispatch(updateSearchMetaData({ tabId, data: { queryId: result.queryId } }));
           startPolling({
             queryId: result.queryId,
           });

public/components/event_analytics/explorer/direct_query_running.tsx

@@ -3,18 +3,42 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { EuiButton, EuiEmptyPrompt, EuiProgress, EuiSpacer, EuiText } from '@elastic/eui';
 import React from 'react';
 import { useDispatch, useSelector } from 'react-redux';
-import { EuiProgress, EuiEmptyPrompt, EuiButton, EuiText, EuiSpacer } from '@elastic/eui';
+import { DirectQueryLoadingStatus } from '../../../../common/types/explorer';
+import { coreRefs } from '../../../framework/core_refs';
+import { SQLService } from '../../../services/requests/sql';
 import {
   selectSearchMetaData,
   update as updateSearchMetaData,
 } from '../redux/slices/search_meta_data_slice';
-import { DirectQueryLoadingStatus } from '../../../../common/types/explorer';
 
 export const DirectQueryRunning = ({ tabId }: { tabId: string }) => {
   const explorerSearchMeta = useSelector(selectSearchMetaData)[tabId] || {};
   const dispatch = useDispatch();
+  const sqlService = new SQLService(coreRefs.http);
+
+  const cancelQuery = () => {
+    if (explorerSearchMeta.queryId !== '') {
+      sqlService
+        .deleteWithJobId({ queryId: explorerSearchMeta.queryId })
+        .catch((e) => {
+          console.error(e);
+        })
+        .finally(() => {
+          dispatch(
+            updateSearchMetaData({
+              tabId,
+              data: {
+                isPolling: false,
+              },
+            })
+          );
+        });
+    }
+  };
+
   return (
     <EuiEmptyPrompt
       icon={<EuiProgress size="xs" color="accent" />}
@@ -25,19 +49,7 @@ export const DirectQueryRunning = ({ tabId }: { tabId: string }) => {
             Status: {explorerSearchMeta.status ?? DirectQueryLoadingStatus.SCHEDULED}
           </EuiText>
           <EuiSpacer size="s" />
-          <EuiButton
-            color="success"
-            onClick={() => {
-              dispatch(
-                updateSearchMetaData({
-                  tabId,
-                  data: {
-                    isPolling: false,
-                  },
-                })
-              );
-            }}
-          >
+          <EuiButton color="success" onClick={cancelQuery}>
             Cancel
           </EuiButton>
         </>

public/components/event_analytics/redux/slices/search_meta_data_slice.ts

@@ -3,14 +3,15 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { createSlice, createSelector, PayloadAction } from '@reduxjs/toolkit';
-import { initialTabId } from '../../../../framework/redux/store/shared_state';
+import { PayloadAction, createSelector, createSlice } from '@reduxjs/toolkit';
 import { REDUX_EXPL_SLICE_SEARCH_META_DATA } from '../../../../../common/constants/explorer';
 import { DirectQueryLoadingStatus, SelectedDataSource } from '../../../../../common/types/explorer';
+import { initialTabId } from '../../../../framework/redux/store/shared_state';
 
 const searchMetaInitialState = {
   lang: 'PPL',
   datasources: [],
+  queryId: '',
   isPolling: false,
 };
 
@@ -24,6 +25,7 @@ interface SearchMetaData {
   lang: string;
   datasources: SelectedDataSource[];
   isPolling: boolean;
+  queryId: string;
   status: DirectQueryLoadingStatus;
 }
 

public/services/requests/sql.ts

@@ -6,6 +6,7 @@
 import { CoreStart } from '../../../../../src/core/public';
 import { DirectQueryRequest } from '../../../common/types/explorer';
 
+
 export class SQLService {
   private http;
   constructor(http: CoreStart['http']) {
@@ -31,4 +32,12 @@ export class SQLService {
       throw error;
     });
   };
+
+  deleteWithJobId = async (params: { queryId: string }, errorHandler?: (error: any) => void) => {
+    return this.http.delete(`/api/observability/query/jobs/${params.queryId}`).catch((error) => {
+      console.error('delete error: ', error.body);
+      if (errorHandler) errorHandler(error);
+      throw error;
+    });
+  };
 }

server/adaptors/opensearch_observability_plugin.ts

@@ -117,6 +117,7 @@ export function OpenSearchObservabilityPlugin(Client: any, config: any, componen
     method: 'DELETE',
   });
 
+  // Get async job status
   observability.getJobStatus = clientAction({
     url: {
       fmt: `${JOBS_ENDPOINT_BASE}/<%=queryId%>`,
@@ -130,6 +131,21 @@ export function OpenSearchObservabilityPlugin(Client: any, config: any, componen
     method: 'GET',
   });
 
+  // Delete async job
+  observability.deleteJob = clientAction({
+    url: {
+      fmt: `${JOBS_ENDPOINT_BASE}/<%=queryId%>`,
+      req: {
+        queryId: {
+          type: 'string',
+          required: true,
+        },
+      },
+    },
+    method: 'DELETE',
+  });
+
+  // Run async job
   observability.runDirectQuery = clientAction({
     url: {
       fmt: `${JOBS_ENDPOINT_BASE}`,

server/routes/datasources/datasources_router.ts

@@ -71,4 +71,33 @@ export function registerDatasourcesRoute(router: IRouter) {
       }
     }
   );
+
+  router.delete(
+    {
+      path: `${OBSERVABILITY_BASE}${JOBS_BASE}/{queryId}`,
+      validate: {
+        params: schema.object({
+          queryId: schema.string(),
+        }),
+      },
+    },
+    async (context, request, response): Promise<any> => {
+      try {
+        const res = await context.observability_plugin.observabilityClient
+          .asScoped(request)
+          .callAsCurrentUser('observability.deleteJob', {
+            queryId: request.params.queryId,
+          });
+        return response.ok({
+          body: res,
+        });
+      } catch (error: any) {
+        console.error('Error in deleting job:', error);
+        return response.custom({
+          statusCode: error.statusCode || 500,
+          body: error.message,
+        });
+      }
+    }
+  );
 }