From 3754d2c32cdced24cad50dfe01cff88bce97bc05 Mon Sep 17 00:00:00 2001
From: Asier Isayas <asier.isayas@gmail.com>
Date: Thu, 5 Oct 2023 10:37:59 -0400
Subject: [PATCH] Cancel Query Option (#1646)

* cancel query option
---------

Co-authored-by: Asier Isayas <aisayas@microsoft.com>
---
 src/Common/ErrorHandlingUtils.ts              |  5 +++++
 src/Explorer/Tabs/DocumentsTab.html           | 12 ++++++++++
 src/Explorer/Tabs/DocumentsTab.ts             |  8 ++++++-
 .../Tabs/QueryTab/QueryTabComponent.tsx       | 22 ++++++++++++++++++-
 4 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/Common/ErrorHandlingUtils.ts b/src/Common/ErrorHandlingUtils.ts
index 8a86bdf32..3a537683c 100644
--- a/src/Common/ErrorHandlingUtils.ts
+++ b/src/Common/ErrorHandlingUtils.ts
@@ -51,6 +51,11 @@ const replaceKnownError = (errorMessage: string): string => {
     return "Database throughput is not supported for internal subscriptions.";
   } else if (errorMessage?.indexOf("Partition key paths must contain only valid") >= 0) {
     return "Partition key paths must contain only valid characters and not contain a trailing slash or wildcard character.";
+  } else if (
+    errorMessage?.indexOf("The user aborted a request") >= 0 ||
+    errorMessage?.indexOf("The operation was aborted") >= 0
+  ) {
+    return "User aborted query.";
   }
 
   return errorMessage;
diff --git a/src/Explorer/Tabs/DocumentsTab.html b/src/Explorer/Tabs/DocumentsTab.html
index 096db482e..094c18e10 100644
--- a/src/Explorer/Tabs/DocumentsTab.html
+++ b/src/Explorer/Tabs/DocumentsTab.html
@@ -106,6 +106,18 @@
               Apply Filter
             </button>
           </span>
+          <span class="filterbuttonpad">
+            <button
+              class="filterbtnstyle queryButton"
+              data-bind="
+                                visible: !isPreferredApiMongoDB && isExecuting,
+                                click: onAbortQueryClick"
+              aria-label="Cancel Query"
+              tabindex="0"
+            >
+              Cancel Query
+            </button>
+          </span>
           <span
             class="filterclose"
             role="button"
diff --git a/src/Explorer/Tabs/DocumentsTab.ts b/src/Explorer/Tabs/DocumentsTab.ts
index 1e0219a7c..cd43ab29c 100644
--- a/src/Explorer/Tabs/DocumentsTab.ts
+++ b/src/Explorer/Tabs/DocumentsTab.ts
@@ -78,6 +78,7 @@ export default class DocumentsTab extends TabsBase {
   private _documentsIterator: QueryIterator<ItemDefinition & Resource>;
   private _resourceTokenPartitionKey: string;
   private _isQueryCopilotSampleContainer: boolean;
+  private queryAbortController: AbortController;
 
   constructor(options: ViewModels.DocumentsTabOptions) {
     super(options);
@@ -401,6 +402,10 @@ export default class DocumentsTab extends TabsBase {
     return true;
   };
 
+  public onAbortQueryClick(): void {
+    this.queryAbortController.abort();
+  }
+
   public onDocumentIdClick(clickedDocumentId: DocumentId): Q.Promise<any> {
     if (this.editorState() !== ViewModels.DocumentExplorerState.noDocumentSelected) {
       return Q();
@@ -688,6 +693,7 @@ export default class DocumentsTab extends TabsBase {
   }
 
   public createIterator(): QueryIterator<ItemDefinition & Resource> {
+    this.queryAbortController = new AbortController();
     const filter: string = this.filterContent().trim();
     const query: string = this.buildQuery(filter);
     let options: any = {};
@@ -696,7 +702,7 @@ export default class DocumentsTab extends TabsBase {
     if (this._resourceTokenPartitionKey) {
       options.partitionKey = this._resourceTokenPartitionKey;
     }
-
+    options.abortSignal = this.queryAbortController.signal;
     return this._isQueryCopilotSampleContainer
       ? querySampleDocuments(query, options)
       : queryDocuments(this.collection.databaseId, this.collection.id(), query, options);
diff --git a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx b/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx
index e7b9bb07a..98384d92e 100644
--- a/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx
+++ b/src/Explorer/Tabs/QueryTab/QueryTabComponent.tsx
@@ -8,6 +8,7 @@ import React, { Fragment } from "react";
 import SplitterLayout from "react-splitter-layout";
 import "react-splitter-layout/lib/index.css";
 import LaunchCopilot from "../../../../images/CopilotTabIcon.svg";
+import CancelQueryIcon from "../../../../images/Entity_cancel.svg";
 import ExecuteQueryIcon from "../../../../images/ExecuteQuery.svg";
 import SaveQueryIcon from "../../../../images/save-cosmos.svg";
 import { NormalizedEventKey, QueryCopilotSampleDatabaseId } from "../../../Common/Constants";
@@ -91,6 +92,7 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
   public isCloseClicked: boolean;
   public isCopilotTabActive: boolean;
   private _iterator: MinimalQueryIterator;
+  private queryAbortController: AbortController;
 
   constructor(props: IQueryTabComponentProps) {
     super(props);
@@ -214,6 +216,7 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
   }
 
   private async _executeQueryDocumentsPage(firstItemIndex: number): Promise<void> {
+    this.queryAbortController = new AbortController();
     if (this._iterator === undefined) {
       this._iterator = this.props.isPreferredApiMongoDB
         ? queryIterator(
@@ -225,7 +228,10 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
             this.props.collection.databaseId,
             this.props.collection.id(),
             this.state.selectedContent || this.state.sqlQueryEditorContent,
-            { enableCrossPartitionQuery: HeadersUtility.shouldEnableCrossPartitionKey() } as FeedOptions,
+            {
+              enableCrossPartitionQuery: HeadersUtility.shouldEnableCrossPartitionKey(),
+              abortSignal: this.queryAbortController.signal,
+            } as FeedOptions,
           );
     }
 
@@ -244,6 +250,7 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
     this.setState({
       isExecuting: true,
     });
+    useCommandBar.getState().setContextButtons(this.getTabsButtons());
 
     try {
       const queryResults: ViewModels.QueryResults = await QueryUtils.queryPagesUntilContentPresent(
@@ -268,6 +275,7 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
         isExecuting: false,
       });
       this.togglesOnFocus();
+      useCommandBar.getState().setContextButtons(this.getTabsButtons());
     }
   }
 
@@ -332,6 +340,18 @@ export default class QueryTabComponent extends React.Component<IQueryTabComponen
       buttons.push(launchCopilotButton);
     }
 
+    if (!this.props.isPreferredApiMongoDB && this.state.isExecuting) {
+      const label = "Cancel query";
+      buttons.push({
+        iconSrc: CancelQueryIcon,
+        iconAlt: label,
+        onCommandClick: () => this.queryAbortController.abort(),
+        commandButtonLabel: label,
+        ariaLabel: label,
+        hasPopup: false,
+      });
+    }
+
     return buttons;
   }