From b7ab03fd889836d6a2bc50183d5af8fa008a3e0a Mon Sep 17 00:00:00 2001
From: Sean Li <lnse@amazon.com>
Date: Tue, 22 Oct 2024 17:41:34 -0700
Subject: [PATCH] adding new error for error state

Signed-off-by: Sean Li <lnse@amazon.com>
---
 src/plugins/query_enhancements/common/index.ts  |  1 +
 .../query_enhancements/server/routes/index.ts   |  8 +++++++-
 .../server/search/sql_async_search_strategy.ts  | 17 +++++++++--------
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/plugins/query_enhancements/common/index.ts b/src/plugins/query_enhancements/common/index.ts
index f21e3c24507f..d12c0b562710 100644
--- a/src/plugins/query_enhancements/common/index.ts
+++ b/src/plugins/query_enhancements/common/index.ts
@@ -6,3 +6,4 @@
 export * from './constants';
 export * from './types';
 export * from './utils';
+export * from './errors';
diff --git a/src/plugins/query_enhancements/server/routes/index.ts b/src/plugins/query_enhancements/server/routes/index.ts
index cb2dbc644189..dd17231adee2 100644
--- a/src/plugins/query_enhancements/server/routes/index.ts
+++ b/src/plugins/query_enhancements/server/routes/index.ts
@@ -85,8 +85,14 @@ export function defineSearchStrategyRouteProvider(logger: Logger, router: IRoute
           const queryRes: IDataFrameResponse = await searchStrategy.search(context, req as any, {});
           return res.ok({ body: { ...queryRes } });
         } catch (err) {
+          let error;
+          try {
+            error = JSON.parse(err.message);
+          } catch (e) {
+            error = err;
+          }
           return res.custom({
-            statusCode: err.name,
+            statusCode: error.status,
             body: err.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
index a48d47ce7826..fb2517d920e0 100644
--- a/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts
+++ b/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts
@@ -14,7 +14,13 @@ import {
   createDataFrame,
 } from '../../../data/common';
 import { Facet } from '../utils';
-import { buildQueryStatusConfig, getFields, handleFacetError, SEARCH_STRATEGY } from '../../common';
+import {
+  buildQueryStatusConfig,
+  getFields,
+  handleFacetError,
+  SEARCH_STRATEGY,
+  PollingResponseError,
+} from '../../common';
 
 export const sqlAsyncSearchStrategyProvider = (
   config$: Observable<SharedGlobalConfig>,
@@ -78,13 +84,8 @@ export const sqlAsyncSearchStrategyProvider = (
               body: dataFrame,
             } as IDataFrameResponse;
           } else if (queryStatus?.toUpperCase() === 'FAILED') {
-            return {
-              type: DATA_FRAME_TYPES.POLLING,
-              status: 'failed',
-              body: {
-                error: new Error(`JOB: ${pollQueryResultsParams.queryId} failed`),
-              },
-            } as IDataFrameResponse;
+            const error = new PollingResponseError(queryStatusResponse.data.error, 500);
+            throw error;
           }
 
           return {