feat: add UI trigger for query progress

common/constants.ts

@@ -45,3 +45,5 @@ export const OPENSEARCH_API = {
 export const UI_SETTINGS = {};
 
 export const ERROR_DETAILS = { GUARDRAILS_TRIGGERED: 'guardrails triggered' };
+
+export const ASYNC_TRIGGER_ID = 'queryenhancements_async_query_interceptor_trigger';

common/types.ts

@@ -5,6 +5,7 @@
 
 import { CoreSetup } from 'opensearch-dashboards/public';
 import { Observable } from 'rxjs';
+import { ASYNC_TRIGGER_ID } from './constants';
 
 export interface FetchDataFrameContext {
   http: CoreSetup['http'];
@@ -13,3 +14,25 @@ export interface FetchDataFrameContext {
 }
 
 export type FetchFunction<T, P = void> = (params?: P) => Observable<T>;
+
+// ref: https://docs.aws.amazon.com/emr/latest/EMR-Serverless-UserGuide/job-states.html
+export enum SparkJobState {
+  SUBMITTED,
+  PENDING,
+  SCHEDULED,
+  RUNNING,
+  FAILED,
+  SUCCESS,
+  CANCELLING,
+  CANCELLED,
+}
+
+export interface AsyncQueryContext {
+  query_id: string;
+  query_status: SparkJobState;
+}
+declare module '../../../src/plugins/ui_actions/public' {
+  export interface TriggerContextMapping {
+    [ASYNC_TRIGGER_ID]: AsyncQueryContext;
+  }
+}

opensearch_dashboards.json

@@ -4,7 +4,7 @@
   "opensearchDashboardsVersion": "opensearchDashboards",
   "server": true,
   "ui": true,
-  "requiredPlugins": ["data"],
+  "requiredPlugins": ["data", "uiActions"],
   "optionalPlugins": ["dataSource", "dataSourceManagement"],
   "requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact", "dataSourceManagement"]
 }

public/plugin.tsx

@@ -4,6 +4,7 @@
  */
 
 import moment from 'moment';
+import { Trigger } from 'src/plugins/ui_actions/public';
 import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '../../../src/core/public';
 import { IStorageWrapper, Storage } from '../../../src/plugins/opensearch_dashboards_utils/public';
 import { ConfigSchema } from '../common/config';
@@ -16,6 +17,9 @@ import {
   QueryEnhancementsPluginStart,
   QueryEnhancementsPluginStartDependencies,
 } from './types';
+import { ASYNC_TRIGGER_ID } from '../common';
+
+export type PublicConfig = Pick<ConfigSchema, 'queryAssist'>;
 import { ConnectionsService, createDataSourceConnectionExtension } from './data_source_connection';
 
 export class QueryEnhancementsPlugin
@@ -37,11 +41,12 @@ export class QueryEnhancementsPlugin
 
   public setup(
     core: CoreSetup<QueryEnhancementsPluginStartDependencies>,
-    { data }: QueryEnhancementsPluginSetupDependencies
+    { data, uiActions }: QueryEnhancementsPluginSetupDependencies
   ): QueryEnhancementsPluginSetup {
     this.connectionsService = new ConnectionsService({
       startServices: core.getStartServices(),
       http: core.http,
+      uiActions,
     });
 
     const pplSearchInterceptor = new PPLSearchInterceptor(
@@ -51,6 +56,7 @@ export class QueryEnhancementsPlugin
         uiSettings: core.uiSettings,
         startServices: core.getStartServices(),
         usageCollector: data.search.usageCollector,
+        uiActions,
       },
       this.connectionsService
     );
@@ -62,6 +68,7 @@ export class QueryEnhancementsPlugin
         uiSettings: core.uiSettings,
         startServices: core.getStartServices(),
         usageCollector: data.search.usageCollector,
+        uiActions,
       },
       this.connectionsService
     );
@@ -73,6 +80,7 @@ export class QueryEnhancementsPlugin
         uiSettings: core.uiSettings,
         startServices: core.getStartServices(),
         usageCollector: data.search.usageCollector,
+        uiActions,
       },
       this.connectionsService
     );
@@ -109,8 +117,7 @@ export class QueryEnhancementsPlugin
           searchBar: {
             showDatePicker: false,
             showFilterBar: false,
-            showDataSetsSelector: false,
-            showDataSourcesSelector: true,
+            showDataSourceSelector: true,
             queryStringInput: { initialValue: 'SELECT * FROM <data_source>' },
           },
           fields: {
@@ -131,8 +138,7 @@ export class QueryEnhancementsPlugin
           searchBar: {
             showDatePicker: false,
             showFilterBar: false,
-            showDataSetsSelector: false,
-            showDataSourcesSelector: true,
+            showDataSourceSelector: true,
             queryStringInput: { initialValue: 'SHOW DATABASES IN ::mys3::' },
           },
           fields: {
@@ -160,6 +166,11 @@ export class QueryEnhancementsPlugin
       },
     });
 
+    const ASYNC_TRIGGER: Trigger<typeof ASYNC_TRIGGER_ID> = {
+      id: ASYNC_TRIGGER_ID,
+    };
+    uiActions.registerTrigger(ASYNC_TRIGGER);
+
     return {};
   }
 

public/search/sql_async_search_interceptor.ts

@@ -1,7 +1,13 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
 import { trimEnd } from 'lodash';
 import { BehaviorSubject, Observable, throwError } from 'rxjs';
 import { i18n } from '@osd/i18n';
 import { concatMap, map } from 'rxjs/operators';
+import { UiActionsStart } from 'src/plugins/ui_actions/public';
 import {
   DATA_FRAME_TYPES,
   DataPublicPluginStart,
@@ -18,9 +24,11 @@ import {
 } from '../../../../src/plugins/data/common';
 import {
   API,
+  ASYNC_TRIGGER_ID,
   DataFramePolling,
   FetchDataFrameContext,
   SEARCH_STRATEGY,
+  SparkJobState,
   fetchDataFrame,
   fetchDataFramePolling,
 } from '../../common';
@@ -32,14 +40,16 @@ export class SQLAsyncSearchInterceptor extends SearchInterceptor {
   protected aggsService!: DataPublicPluginStart['search']['aggs'];
   protected indexPatterns!: DataPublicPluginStart['indexPatterns'];
   protected dataFrame$ = new BehaviorSubject<IDataFrameResponse | undefined>(undefined);
+  protected uiActions: UiActionsStart;
 
   constructor(
     deps: SearchInterceptorDeps,
     private readonly connectionsService: ConnectionsService
   ) {
     super(deps);
+    this.uiActions = deps.uiActions;
 
-    deps.startServices.then(([coreStart, depsStart]) => {
+    deps.startServices.then(([_coreStart, depsStart]) => {
       this.queryService = (depsStart as QueryEnhancementsPluginStartDependencies).data.query;
       this.aggsService = (depsStart as QueryEnhancementsPluginStartDependencies).data.search.aggs;
     });
@@ -48,7 +58,7 @@ export class SQLAsyncSearchInterceptor extends SearchInterceptor {
   protected runSearch(
     request: IOpenSearchDashboardsSearchRequest,
     signal?: AbortSignal,
-    strategy?: string
+    _strategy?: string
   ): Observable<IOpenSearchDashboardsSearchResponse> {
     const { id, ...searchRequest } = request;
     const path = trimEnd(API.SQL_ASYNC_SEARCH);
@@ -77,18 +87,29 @@ export class SQLAsyncSearchInterceptor extends SearchInterceptor {
     };
 
     const onPollingSuccess = (pollingResult: any) => {
-      if (pollingResult && pollingResult.body.meta.status === 'SUCCESS') {
-        return false;
-      }
-      if (pollingResult && pollingResult.body.meta.status === 'FAILED') {
-        const jsError = new Error(pollingResult.data.error.response);
-        this.deps.toasts.addError(jsError, {
-          title: i18n.translate('queryEnhancements.sqlQueryError', {
-            defaultMessage: 'Could not complete the SQL async query',
-          }),
-          toastMessage: pollingResult.data.error.response,
-        });
-        return false;
+      if (pollingResult) {
+        const statusStr: string = (pollingResult.body.meta.status as string).toUpperCase();
+        const status = SparkJobState[statusStr as keyof typeof SparkJobState];
+        switch (status) {
+          case SparkJobState.SUCCESS:
+            return false;
+          case SparkJobState.FAILED:
+            const jsError = new Error(pollingResult.data.error.response);
+            this.deps.toasts.addError(jsError, {
+              title: i18n.translate('queryEnhancements.sqlQueryError', {
+                defaultMessage: 'Could not complete the SQL async query',
+              }),
+              toastMessage: pollingResult.data.error.response,
+            });
+            return false;
+          default:
+            if (request.params?.progress_query_id) {
+              this.uiActions.getTrigger(ASYNC_TRIGGER_ID).exec({
+                query_id: request.params.progress_query_id,
+                query_status: status,
+              });
+            }
+        }
       }
 
       this.deps.toasts.addInfo({

public/types.ts

@@ -5,6 +5,7 @@
 
 import { CoreSetup, CoreStart } from 'opensearch-dashboards/public';
 import { DataPublicPluginSetup, DataPublicPluginStart } from 'src/plugins/data/public';
+import { UiActionsStart } from 'src/plugins/ui_actions/public';
 import { DataSourcePluginStart } from 'src/plugins/data_source/public';
 
 // eslint-disable-next-line @typescript-eslint/no-empty-interface
@@ -15,6 +16,7 @@ export interface QueryEnhancementsPluginStart {}
 
 export interface QueryEnhancementsPluginSetupDependencies {
   data: DataPublicPluginSetup;
+  uiActions: UiActionsStart;
 }
 
 export interface QueryEnhancementsPluginStartDependencies {
@@ -32,5 +34,6 @@ export interface Connection {
 
 export interface ConnectionsServiceDeps {
   http: CoreSetup['http'];
+  uiActions: UiActionsStart;
   startServices: Promise<[CoreStart, QueryEnhancementsPluginStartDependencies, unknown]>;
 }