fix: cancelling query

ydb-platform · Jan 27, 2025 · 61d0258 · 61d0258
1 parent 65e12bc
commit 61d0258
Show file tree

Hide file tree

Showing 5 changed files with 34 additions and 11 deletions.
diff --git a/src/containers/Tenant/Query/CancelQueryButton/CancelQueryButton.tsx b/src/containers/Tenant/Query/CancelQueryButton/CancelQueryButton.tsx
@@ -14,14 +14,16 @@ const b = cn('cancel-query-button');
 interface CancelQueryButtonProps {
     queryId: string;
     tenantName: string;
+    onClick?: VoidFunction;
 }
 
-export function CancelQueryButton({queryId, tenantName}: CancelQueryButtonProps) {
+export function CancelQueryButton({queryId, tenantName, onClick}: CancelQueryButtonProps) {
     const [sendCancelQuery, cancelQueryResponse] = cancelQueryApi.useCancelQueryMutation();
 
     const onStopButtonClick = React.useCallback(() => {
         sendCancelQuery({queryId, database: tenantName});
-    }, [queryId, sendCancelQuery, tenantName]);
+        onClick?.();
+    }, [onClick, queryId, sendCancelQuery, tenantName]);
 
     return (
         <Button

diff --git a/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx b/src/containers/Tenant/Query/QueryEditor/QueryEditor.tsx
@@ -107,12 +107,13 @@ export default function QueryEditor(props: QueryEditorProps) {
         LAST_USED_QUERY_ACTION_KEY,
     );
     const [lastExecutedQueryText, setLastExecutedQueryText] = React.useState<string>('');
+    const isStreamingSupported = useStreamingAvailable();
 
     const [sendQuery] = queryApi.useUseSendQueryMutation();
-
-    const isStreamingSupported = useStreamingAvailable();
     const [streamQuery] = queryApi.useUseStreamQueryMutation();
 
+    const runningQueryRef = React.useRef<{abort: VoidFunction} | null>(null);
+
     React.useEffect(() => {
         if (savedPath !== tenantName) {
             dispatch(setTenantPath(tenantName));
@@ -155,7 +156,7 @@ export default function QueryEditor(props: QueryEditorProps) {
         const queryId = uuidv4();
 
         if (isStreamingSupported) {
-            streamQuery({
+            runningQueryRef.current = streamQuery({
                 actionType: 'execute',
                 query,
                 database: tenantName,
@@ -164,7 +165,7 @@ export default function QueryEditor(props: QueryEditorProps) {
                 queryId,
             });
         } else {
-            sendQuery({
+            runningQueryRef.current = sendQuery({
                 actionType: 'execute',
                 query,
                 database: tenantName,
@@ -199,7 +200,7 @@ export default function QueryEditor(props: QueryEditorProps) {
 
         const queryId = uuidv4();
 
-        sendQuery({
+        runningQueryRef.current = sendQuery({
             actionType: 'explain',
             query: input,
             database: tenantName,
@@ -221,6 +222,12 @@ export default function QueryEditor(props: QueryEditorProps) {
         }
     });
 
+    const handleCancelRunningQuery = React.useCallback(() => {
+        if (runningQueryRef.current) {
+            runningQueryRef.current.abort();
+        }
+    }, []);
+
     const editorWillUnmount = () => {
         window.ydbEditor = undefined;
     };
@@ -388,6 +395,7 @@ export default function QueryEditor(props: QueryEditorProps) {
                         path={path}
                         showPreview={showPreview}
                         queryText={lastExecutedQueryText}
+                        onCancelRunningQuery={handleCancelRunningQuery}
                     />
                 </div>
             </SplitPane>
@@ -407,6 +415,7 @@ interface ResultProps {
     path: string;
     showPreview?: boolean;
     queryText: string;
+    onCancelRunningQuery: VoidFunction;
 }
 function Result({
     resultVisibilityState,
@@ -419,6 +428,7 @@ function Result({
     path,
     showPreview,
     queryText,
+    onCancelRunningQuery,
 }: ResultProps) {
     if (showPreview) {
         return <Preview database={tenantName} path={path} type={type} />;
@@ -435,6 +445,7 @@ function Result({
                 onExpandResults={onExpandResultHandler}
                 onCollapseResults={onCollapseResultHandler}
                 queryText={queryText}
+                onCancelRunningQuery={onCancelRunningQuery}
             />
         );
     }

diff --git a/src/containers/Tenant/Query/QueryResult/QueryResultViewer.tsx b/src/containers/Tenant/Query/QueryResult/QueryResultViewer.tsx
@@ -81,9 +81,10 @@ interface ExecuteResultProps {
     isResultsCollapsed?: boolean;
     theme?: string;
     tenantName: string;
+    queryText?: string;
+    onCancelRunningQuery?: VoidFunction;
     onCollapseResults: VoidFunction;
     onExpandResults: VoidFunction;
-    queryText?: string;
 }
 
 export function QueryResultViewer({
@@ -93,6 +94,7 @@ export function QueryResultViewer({
     theme,
     tenantName,
     queryText,
+    onCancelRunningQuery,
     onCollapseResults,
     onExpandResults,
 }: ExecuteResultProps) {
@@ -288,7 +290,11 @@ export function QueryResultViewer({
                 {isLoading ? (
                     <React.Fragment>
                         <ElapsedTime className={b('elapsed-time')} />
-                        <CancelQueryButton queryId={queryId} tenantName={tenantName} />
+                        <CancelQueryButton
+                            queryId={queryId}
+                            tenantName={tenantName}
+                            onClick={onCancelRunningQuery}
+                        />
                     </React.Fragment>
                 ) : null}
                 {data?.traceId && isExecute ? (

diff --git a/src/containers/Tenant/Query/QueryResult/components/QueryResultError/QueryResultError.tsx b/src/containers/Tenant/Query/QueryResult/components/QueryResultError/QueryResultError.tsx
@@ -10,7 +10,7 @@ const b = cn('ydb-query-result-error ');
 export function QueryResultError({error}: {error: unknown}) {
     const parsedError = parseQueryError(error);
 
-    // "Stopped" message is displayd in QueryExecutionStatus
+    // "Stopped" message is displayed in QueryExecutionStatus
     // There is no need to display "Query is cancelled" message too
     if (!parsedError || isQueryCancelledError(error)) {
         return null;

diff --git a/src/containers/Tenant/Query/utils/isQueryCancelledError.ts b/src/containers/Tenant/Query/utils/isQueryCancelledError.ts
@@ -1,6 +1,10 @@
+import type {IResponseError} from '../../../../types/api/error';
 import {parseQueryError} from '../../../../utils/query';
 
 export function isQueryCancelledError(error: unknown) {
     const parsedError = parseQueryError(error);
-    return typeof parsedError === 'object' && parsedError.error?.message === 'Query was cancelled';
+    return (
+        (error as IResponseError)?.isCancelled ||
+        (typeof parsedError === 'object' && parsedError.error?.message === 'Query was cancelled')
+    );
 }