Refactor async calls and minor bug fixes

common/types/index.ts

@@ -86,12 +86,14 @@
   formErrors: FormErrorsType;
 }
 
-export enum AsyncQueryLoadingStatus {
+export enum AsyncQueryStatus {
+  Pending = 'pending',
   Success = 'success',
   Failed = 'failed',
   Running = 'running',
   Scheduled = 'scheduled',
   Cancelled = 'cancelled',
+  Waiting = 'waiting',
 }
 
 export type TreeItemType =
@@ -115,3 +117,21 @@
   flag: boolean;
   status: string;
 }
+
+interface AsyncApiDataResponse {
+  status: string;
+  schema?: Array<{ name: string; type: string }>;
+  datarows?: any;
+  total?: number;
+  size?: number;
+  error?: string;
+}
+
+export interface AsyncApiResponse {
+  data: {
+    ok: boolean;
+    resp: AsyncApiDataResponse;
+  };
+}
+
+export type PollingCallback = (statusObj: AsyncApiResponse) => void;

common/utils/async_query_helpers.ts

@@ -3,14 +3,15 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import _ from 'lodash';
-import { CoreStart } from '../../../../src/core/public';
+import { coreRefs } from '../../public/framework/core_refs';
 import {
   ASYNC_QUERY_ENDPOINT,
   ASYNC_QUERY_JOB_ENDPOINT,
   ASYNC_QUERY_SESSION_ID,
   POLL_INTERVAL_MS,
 } from '../constants';
+import { AsyncApiResponse, AsyncQueryStatus, PollingCallback } from '../types';
+import { RaiseErrorToast } from './toast_helper';
 
 export const setAsyncSessionId = (dataSource: string, value: string | null) => {
   if (value !== null) {
@@ -22,55 +23,130 @@ export const getAsyncSessionId = (dataSource: string) => {
   return sessionStorage.getItem(`${ASYNC_QUERY_SESSION_ID}_${dataSource}`);
 };
 
-export const getJobId = (
+export const executeAsyncQuery = (
   currentDataSource: string,
   query: {},
-  http: CoreStart['http'],
-  callback
+  pollingCallback: PollingCallback,
+  onErrorCallback?: (errorMessage: string) => void
 ) => {
-  http
-    .post(ASYNC_QUERY_ENDPOINT, {
-      body: JSON.stringify({
-        ...query,
-        sessionId: getAsyncSessionId(currentDataSource) ?? undefined,
-      }),
-    })
-    .then((res) => {
-      const id = res.data.resp.queryId;
-      setAsyncSessionId(currentDataSource, _.get(res.data.resp, 'sessionId', null));
-      if (id === undefined) {
-        console.error(JSON.parse(res.data.body));
-      }
-      callback(id);
-    })
-    .catch((err) => {
-      console.error(err);
-    });
-};
+  let jobId: string | undefined;
+  let isQueryFulfilled = false;
+  const http = coreRefs.http!;
+
+  const getJobId = () => {
+    http
+      .post(ASYNC_QUERY_ENDPOINT, {
+        body: JSON.stringify({
+          ...query,
+          sessionId: getAsyncSessionId(currentDataSource) ?? undefined,
+        }),
+      })
+      .then((res) => {
+        const responseData = res.data.resp;
+        jobId = responseData?.queryId;
+        if (jobId === undefined || !res.data.ok) {
+          console.error('Recieved an invalid query id:', res.data);
+          RaiseErrorToast({
+            errorToastMessage: 'Recieved an invalid query id: ' + res.data.resp,
+            errorDetailsMessage: res.data.body,
+          });
+
+          if (onErrorCallback) {
+            onErrorCallback(res.data.body);
+          }
+          return;
+        }
+        setAsyncSessionId(currentDataSource, responseData?.sessionId ?? null);
+        pollQueryStatus(jobId, pollingCallback);
+      })
+      .catch((err) => {
+        console.error('Error occurred while getting query id:', err);
+        RaiseErrorToast({
+          errorToastMessage: 'Error occurred while getting query id',
+          errorDetailsMessage: err,
+        });
+        isQueryFulfilled = true;
+        if (onErrorCallback) {
+          onErrorCallback(err);
+        }
+      });
+  };
+
+  const pollQueryStatus = (id: string, callback: PollingCallback) => {
+    http
+      .get(ASYNC_QUERY_JOB_ENDPOINT + id)
+      .then((res: AsyncApiResponse) => {
+        const status = res.data.resp.status.toLowerCase();
+        const errorDetailsMessage = res.data.resp.error ?? '';
+        switch (status) {
+          case AsyncQueryStatus.Pending:
+          case AsyncQueryStatus.Running:
+          case AsyncQueryStatus.Scheduled:
+          case AsyncQueryStatus.Waiting:
+            callback({ ...res });
+            setTimeout(() => pollQueryStatus(id, callback), POLL_INTERVAL_MS);
+            break;
+
+          case AsyncQueryStatus.Failed:
+          case AsyncQueryStatus.Cancelled:
+            isQueryFulfilled = true;
+
+            if (status === AsyncQueryStatus.Failed) {
+              RaiseErrorToast({
+                errorToastMessage: 'Query failed',
+                errorDetailsMessage,
+              });
+            }
+            if (onErrorCallback) {
+              onErrorCallback(errorDetailsMessage);
+            }
+            callback({ ...res });
+            break;
+
+          case AsyncQueryStatus.Success:
+            isQueryFulfilled = true;
+            callback({ ...res });
+            break;
+
+          default:
+            console.error('Unrecognized status:', status);
+            RaiseErrorToast({
+              errorToastMessage: 'Unrecognized status recieved',
+              errorDetailsMessage: 'Unrecognized status recieved - ' + errorDetailsMessage,
+            });
+            if (onErrorCallback) {
+              onErrorCallback(errorDetailsMessage);
+            }
+            callback({ ...res });
+        }
+      })
+      .catch((err) => {
+        console.error('Error occurred while polling query status:', err);
+        isQueryFulfilled = true;
+        callback({
+          data: {
+            ok: true,
+            resp: { status: AsyncQueryStatus.Failed, error: 'Failed to query status' },
+          },
+        });
+      });
+  };
+
+  const cancelQuery = () => {
+    if (jobId && !isQueryFulfilled) {
+      http.delete(ASYNC_QUERY_JOB_ENDPOINT + jobId).catch((err) => {
+        console.error('Error occurred while cancelling query:', err);
+        RaiseErrorToast({
+          errorToastMessage: 'Query cancellation failed',
+          errorDetailsMessage: 'Query cancellation failed for queryId: ' + err,
+        });
+      });
+    }
+  };
+
+  // Start executing the query
+  getJobId();
 
-export const pollQueryStatus = (id: string, http: CoreStart['http'], callback) => {
-  http
-    .get(ASYNC_QUERY_JOB_ENDPOINT + id)
-    .then((res) => {
-      const status = res.data.resp.status.toLowerCase();
-      if (
-        status === 'pending' ||
-        status === 'running' ||
-        status === 'scheduled' ||
-        status === 'waiting'
-      ) {
-        callback({ status });
-        setTimeout(() => pollQueryStatus(id, http, callback), POLL_INTERVAL_MS);
-      } else if (status === 'failed') {
-        const results = res.data.resp;
-        callback({ status: 'FAILED', error: results.error });
-      } else if (status === 'success') {
-        const results = _.get(res.data.resp, 'datarows');
-        callback({ status: 'SUCCESS', results });
-      }
-    })
-    .catch((err) => {
-      console.error(err);
-      callback({ status: 'FAILED', error: 'Failed to fetch data' });
-    });
+  // Return a function to cancel the query
+  return cancelQuery;
 };

common/utils/toast_helper.tsx

@@ -0,0 +1,90 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+  EuiButton,
+  EuiCodeBlock,
+  EuiModal,
+  EuiModalBody,
+  EuiModalFooter,
+  EuiModalHeader,
+  EuiModalHeaderTitle,
+} from '@elastic/eui';
+import React from 'react';
+import { MountPoint, ToastInputFields } from '../../../../src/core/public';
+import { toMountPoint } from '../../../../src/plugins/opensearch_dashboards_react/public';
+import { coreRefs } from '../../public/framework/core_refs';
+
+type Color = 'success' | 'primary' | 'warning' | 'danger' | undefined;
+
+export const useToast = () => {
+  const toasts = coreRefs.toasts!;
+
+  const setToast = (title: string, color: Color = 'success', text?: string | MountPoint) => {
+    const newToast: ToastInputFields = {
+      id: new Date().toISOString(),
+      title,
+      text,
+    };
+    switch (color) {
+      case 'danger': {
+        toasts.addDanger(newToast);
+        break;
+      }
+      case 'warning': {
+        toasts.addWarning(newToast);
+        break;
+      }
+      default: {
+        toasts.addSuccess(newToast);
+        break;
+      }
+    }
+  };
+
+  return { setToast };
+};
+
+const loadErrorModal = (errorDetailsMessage: string) => {
+  const openModal = coreRefs.overlays?.openModal!;
+  const modal = openModal(
+    toMountPoint(
+      <EuiModal onClose={() => modal.close()}>
+        <EuiModalHeader>
+          <EuiModalHeaderTitle>Error details</EuiModalHeaderTitle>
+        </EuiModalHeader>
+        <EuiModalBody>
+          <EuiCodeBlock language="json" fontSize="m" paddingSize="m" isCopyable>
+            {errorDetailsMessage}
+          </EuiCodeBlock>
+        </EuiModalBody>
+        <EuiModalFooter>
+          <EuiButton onClick={() => modal.close()}>Close</EuiButton>
+        </EuiModalFooter>
+      </EuiModal>
+    )
+  );
+};
+
+export const RaiseErrorToast = ({
+  errorToastMessage,
+  errorDetailsMessage,
+}: {
+  errorToastMessage: string;
+  errorDetailsMessage: string;
+}) => {
+  const { setToast } = useToast();
+  setToast(
+    errorToastMessage,
+    'danger',
+    toMountPoint(
+      <EuiButton color="danger" size="s" onClick={() => loadErrorModal(errorDetailsMessage)}>
+        Error details
+      </EuiButton>
+    )
+  );
+
+  return <></>;
+};