[Discover] fix error handling for sql and jobs APIs (#8765) (#8770)

(cherry picked from commit c9941f7) Signed-off-by: Joshua Li <[email protected]> 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>
opensearch-project · Oct 31, 2024 · e2aaab4 · e2aaab4
1 parent 3ddd89b
commit e2aaab4
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 28 deletions.
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
@@ -52,14 +52,20 @@ export function registerDataSourceConnectionsRoutes(
       },
     },
     async (context, request, response) => {
-      const client = request.query.id
-        ? context.dataSource.opensearch.legacy.getClient(request.query.id).callAPI
-        : defaultClient.asScoped(request).callAsCurrentUser;
+      try {
+        const client = request.query.id
+          ? context.dataSource.opensearch.legacy.getClient(request.query.id).callAPI
+          : defaultClient.asScoped(request).callAsCurrentUser;
 
-      const resp = await client('enhancements.getJobStatus', {
-        queryId: request.query.queryId,
-      });
-      return response.ok({ body: resp });
+        const resp = await client('enhancements.getJobStatus', {
+          queryId: request.query.queryId,
+        });
+        return response.ok({ body: resp });
+      } catch (error) {
+        // Transform 500 errors to 503 to indicate service availability issues
+        const statusCode = error.statusCode === 500 ? 503 : error.statusCode || 503;
+        return response.custom({ statusCode, body: error.message });
+      }
     }
   );
 
@@ -79,12 +85,18 @@ export function registerDataSourceConnectionsRoutes(
       },
     },
     async (context, request, response) => {
-      const client = request.query.id
-        ? context.dataSource.opensearch.legacy.getClient(request.query.id).callAPI
-        : defaultClient.asScoped(request).callAsCurrentUser;
+      try {
+        const client = request.query.id
+          ? context.dataSource.opensearch.legacy.getClient(request.query.id).callAPI
+          : defaultClient.asScoped(request).callAsCurrentUser;
 
-      const resp = await client('enhancements.runDirectQuery', { body: request.body });
-      return response.ok({ body: resp });
+        const resp = await client('enhancements.runDirectQuery', { body: request.body });
+        return response.ok({ body: resp });
+      } catch (error) {
+        // Transform 500 errors to 503 to indicate service availability issues
+        const statusCode = error.statusCode === 500 ? 503 : error.statusCode || 503;
+        return response.custom({ statusCode, body: error.message });
+      }
     }
   );
 }
diff --git a/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts b/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts
@@ -59,6 +59,9 @@ export const sqlAsyncSearchStrategyProvider = (
         } else {
           request.params = { queryId: inProgressQueryId };
           const queryStatusResponse: any = await sqlAsyncJobsFacet.describeQuery(context, request);
+
+          if (!queryStatusResponse.success) handleFacetError(queryStatusResponse);
+
           const queryStatus = queryStatusResponse?.data?.status;
           logger.info(`sqlAsyncSearchStrategy: JOB: ${inProgressQueryId} - STATUS: ${queryStatus}`);
 

diff --git a/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts b/src/plugins/query_enhancements/server/search/sql_search_strategy.test.ts
@@ -3,24 +3,21 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { sqlSearchStrategyProvider } from './sql_search_strategy';
-import { Observable, of } from 'rxjs';
 import {
-  SharedGlobalConfig,
-  Logger,
   ILegacyClusterClient,
+  Logger,
   RequestHandlerContext,
+  SharedGlobalConfig,
 } from 'opensearch-dashboards/server';
+import { Observable, of } from 'rxjs';
+import { DATA_FRAME_TYPES, IOpenSearchDashboardsSearchRequest } from '../../../data/common';
 import { SearchUsage } from '../../../data/server';
-import {
-  DATA_FRAME_TYPES,
-  IDataFrameError,
-  IOpenSearchDashboardsSearchRequest,
-} from '../../../data/common';
-import * as facet from '../utils/facet';
 import * as utils from '../../common/utils';
+import * as facet from '../utils/facet';
+import { sqlSearchStrategyProvider } from './sql_search_strategy';
 
 jest.mock('../../common/utils', () => ({
+  ...jest.requireActual('../../common/utils'),
   getFields: jest.fn(),
 }));
 
@@ -145,6 +142,27 @@ describe('sqlSearchStrategyProvider', () => {
     expect(usage.trackError).toHaveBeenCalled();
   });
 
+  it('should throw error when describeQuery success is false', async () => {
+    const mockError = new Error('Something went wrong');
+    const mockFacet = ({
+      describeQuery: jest.fn().mockResolvedValue({ success: false, data: mockError }),
+    } as unknown) as facet.Facet;
+    jest.spyOn(facet, 'Facet').mockImplementation(() => mockFacet);
+
+    const strategy = sqlSearchStrategyProvider(config$, logger, client, usage);
+    await expect(
+      strategy.search(
+        emptyRequestHandlerContext,
+        ({
+          body: { query: { query: 'SELECT * FROM table' } },
+        } as unknown) as IOpenSearchDashboardsSearchRequest<unknown>,
+        {}
+      )
+    ).rejects.toThrowError();
+    expect(logger.error).toHaveBeenCalledWith(expect.stringContaining(mockError.message));
+    expect(usage.trackError).toHaveBeenCalled();
+  });
+
   it('should handle empty search response', async () => {
     const mockResponse = {
       success: true,

diff --git a/src/plugins/query_enhancements/server/search/sql_search_strategy.ts b/src/plugins/query_enhancements/server/search/sql_search_strategy.ts
@@ -13,7 +13,7 @@ import {
   Query,
   createDataFrame,
 } from '../../../data/common';
-import { getFields } from '../../common/utils';
+import { getFields, handleFacetError } from '../../common/utils';
 import { Facet } from '../utils';
 
 export const sqlSearchStrategyProvider = (
@@ -36,11 +36,7 @@ export const sqlSearchStrategyProvider = (
         const query: Query = request.body.query;
         const rawResponse: any = await sqlFacet.describeQuery(context, request);
 
-        if (!rawResponse.success) {
-          const error = new Error(rawResponse.data.body);
-          error.name = rawResponse.data.status;
-          throw error;
-        }
+        if (!rawResponse.success) handleFacetError(rawResponse);
 
         const dataFrame = createDataFrame({
           name: query.dataset?.id,

diff --git a/src/plugins/query_enhancements/server/utils/facet.test.ts b/src/plugins/query_enhancements/server/utils/facet.test.ts
@@ -8,6 +8,7 @@ import { Facet, FacetProps } from './facet';
 
 describe('Facet', () => {
   let facet: Facet;
+  let facetWithShimEnabled: Facet;
   let mockClient: jest.Mock;
   let mockLogger: jest.Mocked<Logger>;
   let mockContext: any;
@@ -26,6 +27,7 @@ describe('Facet', () => {
     };
 
     facet = new Facet(props);
+    facetWithShimEnabled = new Facet({ ...props, shimResponse: true });
 
     mockContext = {
       dataSource: {
@@ -115,5 +117,17 @@ describe('Facet', () => {
         'Facet fetch: test-endpoint: Error: Test error'
       );
     });
+
+    it('should handle errors with shim enabled', async () => {
+      const error = new Error('Test error');
+      mockClient.mockRejectedValue(error);
+
+      const result = await facetWithShimEnabled.describeQuery(mockContext, mockRequest);
+
+      expect(result).toEqual({ success: false, data: error });
+      expect(mockLogger.error).toHaveBeenCalledWith(
+        'Facet fetch: test-endpoint: Error: Test error'
+      );
+    });
   });
 });