From 7c303c24fc76f24832aa1f1b58039cc22405268d Mon Sep 17 00:00:00 2001
From: Shenoy Pratik <sgguruda@amazon.com>
Date: Mon, 25 Mar 2024 15:11:50 -0700
Subject: [PATCH] acceleration create button update

Signed-off-by: Shenoy Pratik <sgguruda@amazon.com>
---
 .../create/create_acceleration.tsx                            | 1 +
 .../create/create_acceleration_button.tsx                     | 4 +++-
 public/framework/datasources/direct_query_hook.tsx            | 1 +
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration.tsx b/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration.tsx
index 9957d16dd..da0dc8707 100644
--- a/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration.tsx
+++ b/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration.tsx
@@ -44,6 +44,7 @@ export interface CreateAccelerationProps {
   resetFlyout: () => void;
   databaseName?: string;
   tableName?: string;
+  refreshHandler?: () => void;
 }
 
 export const CreateAcceleration = ({
diff --git a/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration_button.tsx b/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration_button.tsx
index a1cbd3a09..8a8724827 100644
--- a/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration_button.tsx
+++ b/public/components/datasources/components/manage/accelerations/create_accelerations_flyout/create/create_acceleration_button.tsx
@@ -62,9 +62,11 @@ export const CreateAccelerationButton = ({
     }
   }, [directqueryLoadStatus]);
 
-  return (
+  const createAccelerationBtn = (
     <EuiButton onClick={createAcceleration} fill isLoading={isLoading}>
       {isLoading ? 'Creating acceleration' : 'Create acceleration'}
     </EuiButton>
   );
+
+  return createAccelerationBtn;
 };
diff --git a/public/framework/datasources/direct_query_hook.tsx b/public/framework/datasources/direct_query_hook.tsx
index f4776385e..a896a5395 100644
--- a/public/framework/datasources/direct_query_hook.tsx
+++ b/public/framework/datasources/direct_query_hook.tsx
@@ -74,6 +74,7 @@ export const useDirectQuery = () => {
       setLoadStatus(status);
       stopLoading();
     } else if (status === DirectQueryLoadingStatus.FAILED) {
+      setLoadStatus(status);
       stopLoading();
       const formattedError = formatError(
         '',