[Discover] Cancel S3 Queries Using SQL Plugin (#9355)

opensearch-project · Feb 19, 2025 · 3c90fa3 · 3c90fa3
1 parent 3918797
commit 3c90fa3
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 7 deletions.
diff --git a/changelogs/fragments/9355.yml b/changelogs/fragments/9355.yml
@@ -0,0 +1,2 @@
+feat:
+- Deletes S3 Jobs in Backend when Original Query is Canceled ([#9355](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/9355))
diff --git a/...re_opensearch_dashboards/opensearch_dashboards/apps/query_enhancements/s3_dataset.spec.js b/...re_opensearch_dashboards/opensearch_dashboards/apps/query_enhancements/s3_dataset.spec.js
@@ -7,6 +7,7 @@ import {
   DS_API,
   DSM_API,
   S3_CLUSTER,
+  JOBS_API,
 } from '../../../../../utils/apps/query_enhancements/constants';
 import { getRandomizedWorkspaceName } from '../../../../../utils/apps/query_enhancements/shared';
 import { prepareTestSuite } from '../../../../../utils/helpers';
@@ -90,6 +91,11 @@ const s3DatasetTestSuite = () => {
           cy.deleteWorkspaceByName(workspace);
           cy.visit('/app/home');
           cy.osd.createInitialWorkspaceWithDataSource(S3_CLUSTER.name, workspace);
+          cy.navigateToWorkSpaceSpecificPage({
+            workspaceName: workspace,
+            page: 'discover',
+            isEnhancement: true,
+          });
         });
         afterEach(() => {
           cy.deleteWorkspaceByName(workspace);
@@ -141,6 +147,40 @@ const s3DatasetTestSuite = () => {
           cy.getElementByTestId('docTable').should('be.visible');
           cy.getElementByTestId('docTable').find('tr').should('have.length', 11);
         });
+
+        it('aborts and cancels previous query when new query is started', function () {
+          cy.getElementByTestId(`datasetSelectorButton`).click();
+          cy.getElementByTestId(`datasetSelectorAdvancedButton`).click();
+
+          cy.get(`[title="S3 Connections"]`).click();
+          cy.get(`[title="BasicS3Connection"]`).click();
+          cy.get(`[title="mys3"]`).click();
+          cy.get(`[title="default"]`).click();
+          cy.get(`[title="http_logs"]`).click();
+          cy.getElementByTestId('datasetSelectorNext').click();
+          cy.get(`[class="euiModalHeader__title"]`).should('contain', 'Step 2: Configure data');
+
+          cy.getElementByTestId('advancedSelectorLanguageSelect').select('OpenSearch SQL');
+          cy.getElementByTestId('advancedSelectorConfirmButton').click();
+
+          // Need to wait a bit for initial query to start
+          cy.wait(3000);
+
+          cy.intercept('DELETE', `**/${JOBS_API.DELETE}*`).as('cancelRequest');
+          cy.getElementByTestId(`querySubmitButton`).click();
+
+          cy.wait('@cancelRequest').then((interception) => {
+            console.log('interception.request.url:', interception.request.url);
+            // Verify the request had the correct query parameters
+            expect(interception.request.url).to.include('queryId=');
+            expect(interception.request.url).to.include('id=');
+          });
+
+          cy.waitForSearch();
+          cy.get(`[data-test-subj="queryResultCompleteMsg"]`).should('be.visible');
+          cy.getElementByTestId('docTable').should('be.visible');
+          cy.getElementByTestId('docTable').find('tr').should('have.length', 11);
+        });
       });
     }
   );

diff --git a/cypress/utils/apps/query_enhancements/constants.js b/cypress/utils/apps/query_enhancements/constants.js
@@ -21,6 +21,11 @@ export const DS_API = {
 };
 export const DSM_API = '/internal/data-source-management/fetchDataSourceMetaData';
 
+export const BASE_QUERY_ENHANCEMENTS_API = '/api/enhancements';
+export const JOBS_API = {
+  DELETE: `${BASE_QUERY_ENHANCEMENTS_API}/jobs`,
+};
+
 export const INDEX_WITH_TIME_1 = 'data_logs_small_time_1';
 export const INDEX_WITHOUT_TIME_1 = 'data_logs_small_no_time_1';
 export const INDEX_WITH_TIME_2 = 'data_logs_small_time_2';

diff --git a/src/plugins/query_enhancements/common/utils.ts b/src/plugins/query_enhancements/common/utils.ts
@@ -13,6 +13,7 @@ import {
   QueryStatusConfig,
   QueryStatusOptions,
 } from './types';
+import { API } from './constants';
 
 export const formatDate = (dateString: string) => {
   const date = new Date(dateString);
@@ -57,13 +58,34 @@ export const fetch = (context: EnhancedFetchContext, query: Query, aggConfig?: Q
     pollQueryResultsParams: context.body?.pollQueryResultsParams,
     timeRange: context.body?.timeRange,
   });
+
   return from(
-    http.fetch({
-      method: 'POST',
-      path,
-      body,
-      signal,
-    })
+    http
+      .fetch({
+        method: 'POST',
+        path,
+        body,
+        signal,
+      })
+      .catch(async (error) => {
+        if (error.name === 'AbortError' && context.body?.pollQueryResultsParams?.queryId) {
+          // Cancel job
+          try {
+            await http.fetch({
+              method: 'DELETE',
+              path: API.DATA_SOURCE.ASYNC_JOBS,
+              query: {
+                id: query.dataset?.dataSource?.id,
+                queryId: context.body?.pollQueryResultsParams.queryId,
+              },
+            });
+          } catch (cancelError) {
+            // eslint-disable-next-line no-console
+            console.error('Failed to cancel query:', cancelError);
+          }
+          throw error;
+        }
+      })
   );
 };
 

diff --git a/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts b/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts
@@ -99,4 +99,31 @@ export function registerDataSourceConnectionsRoutes(
       }
     }
   );
+
+  router.delete(
+    {
+      path: API.DATA_SOURCE.ASYNC_JOBS,
+      validate: {
+        query: schema.object({
+          id: schema.string(),
+          queryId: schema.nullable(schema.string()),
+        }),
+      },
+    },
+    async (context, request, response) => {
+      try {
+        const client = request.query.id
+          ? context.dataSource.opensearch.legacy.getClient(request.query.id).callAPI
+          : defaultClient.asScoped(request).callAsCurrentUser;
+
+        await client('enhancements.deleteJob', {
+          queryId: request.query.queryId,
+        });
+        return response.noContent();
+      } catch (error) {
+        const statusCode = error.statusCode === 500 ? 503 : error.statusCode || 503;
+        return response.custom({ statusCode, body: error.message });
+      }
+    }
+  );
 }
diff --git a/src/plugins/query_enhancements/server/utils/plugins.ts b/src/plugins/query_enhancements/server/utils/plugins.ts
@@ -141,7 +141,7 @@ export const OpenSearchEnhancements = (client: any, config: any, components: any
   });
 
   enhancements.deleteJob = createAction(client, components, {
-    endpoint: `${URI.ASYNC_QUERY}/<%=queryId%>`,
+    endpoint: `${URI.ASYNC_QUERY}`,
     method: 'DELETE',
     paramKey: 'queryId',
   });
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		feat:
		- Deletes S3 Jobs in Backend when Original Query is Canceled ([#9355](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/9355))