From da2e874ae6b5a6dab613a306add65f297250f69a Mon Sep 17 00:00:00 2001
From: sunghyunkang1111 <114709653+sunghyunkang1111@users.noreply.github.com>
Date: Thu, 21 Mar 2024 11:23:42 -0500
Subject: [PATCH] Fix bugs on data transfer and bring back query explanation
 and remove query prompt from editor (#1777)

* Fix minor issues

* add back preview tag

* bring back query explanation and remove prompt in editor
---
 src/Common/dataAccess/dataTransfers.ts             | 11 +++++++++--
 .../PartitionKeyComponent.tsx                      | 14 +++++++-------
 src/Explorer/Controls/Settings/SettingsUtils.tsx   |  2 +-
 .../__snapshots__/SettingsComponent.test.tsx.snap  |  2 +-
 .../ChangePartitionKeyPane.tsx                     |  2 ++
 .../QueryCopilot/QueryCopilotPromptbar.tsx         |  6 ++++--
 6 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/src/Common/dataAccess/dataTransfers.ts b/src/Common/dataAccess/dataTransfers.ts
index 138257c15..e639f9965 100644
--- a/src/Common/dataAccess/dataTransfers.ts
+++ b/src/Common/dataAccess/dataTransfers.ts
@@ -122,14 +122,21 @@ const pollDataTransferJobOperation = async (
 
   updateDataTransferJob(body);
 
-  if (status === "Cancelled" || status === "Failed" || status === "Faulted") {
+  if (status === "Cancelled") {
+    removeFromPolling(jobName);
+    clearMessage && clearMessage();
+    const cancelMessage = `Data transfer job ${jobName} cancelled`;
+    NotificationConsoleUtils.logConsoleError(cancelMessage);
+    throw new AbortError(cancelMessage);
+  }
+  if (status === "Failed" || status === "Faulted") {
     removeFromPolling(jobName);
     const errorMessage = body?.properties?.error
       ? JSON.stringify(body?.properties?.error)
       : "Operation could not be completed";
     const error = new Error(errorMessage);
     clearMessage && clearMessage();
-    NotificationConsoleUtils.logConsoleError(`Data transfer job ${jobName} Failed`);
+    NotificationConsoleUtils.logConsoleError(`Data transfer job ${jobName} failed: ${errorMessage}`);
     throw new AbortError(error);
   }
   if (status === "Completed") {
diff --git a/src/Explorer/Controls/Settings/SettingsSubComponents/PartitionKeyComponent.tsx b/src/Explorer/Controls/Settings/SettingsSubComponents/PartitionKeyComponent.tsx
index 2efe7fb92..1930808a7 100644
--- a/src/Explorer/Controls/Settings/SettingsSubComponents/PartitionKeyComponent.tsx
+++ b/src/Explorer/Controls/Settings/SettingsSubComponents/PartitionKeyComponent.tsx
@@ -136,15 +136,15 @@ export const PartitionKeyComponent: React.FC<PartitionKeyComponentProps> = ({ da
   };
 
   const getPercentageComplete = () => {
+    const jobStatus = portalDataTransferJob?.properties?.status;
+    const isCompleted = jobStatus === "Completed";
+    if (isCompleted) {
+      return 1;
+    }
     const processedCount = portalDataTransferJob?.properties?.processedCount;
     const totalCount = portalDataTransferJob?.properties?.totalCount;
-    const jobStatus = portalDataTransferJob?.properties?.status;
-    const isCancelled = jobStatus === "Cancelled";
-    const isCompleted = jobStatus === "Completed";
-    if (totalCount <= 0 && !isCompleted) {
-      return isCancelled ? 0 : null;
-    }
-    return isCompleted ? 1 : processedCount / totalCount;
+    const isJobInProgress = isCurrentJobInProgress(portalDataTransferJob);
+    return isJobInProgress ? (totalCount > 0 ? processedCount / totalCount : null) : 0;
   };
 
   return (
diff --git a/src/Explorer/Controls/Settings/SettingsUtils.tsx b/src/Explorer/Controls/Settings/SettingsUtils.tsx
index c53974b88..869ae323a 100644
--- a/src/Explorer/Controls/Settings/SettingsUtils.tsx
+++ b/src/Explorer/Controls/Settings/SettingsUtils.tsx
@@ -149,7 +149,7 @@ export const getTabTitle = (tab: SettingsV2TabTypes): string => {
     case SettingsV2TabTypes.IndexingPolicyTab:
       return "Indexing Policy";
     case SettingsV2TabTypes.PartitionKeyTab:
-      return "Partition Keys";
+      return "Partition Keys (preview)";
     case SettingsV2TabTypes.ComputedPropertiesTab:
       return "Computed Properties (preview)";
     default:
diff --git a/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap b/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap
index 3b9f76195..ab7abac11 100644
--- a/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap
+++ b/src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap
@@ -207,7 +207,7 @@ exports[`SettingsComponent renders 1`] = `
         />
       </PivotItem>
       <PivotItem
-        headerText="Partition Keys"
+        headerText="Partition Keys (preview)"
         itemKey="PartitionKeyTab"
         key="PartitionKeyTab"
         style={
diff --git a/src/Explorer/Panes/ChangePartitionKeyPane/ChangePartitionKeyPane.tsx b/src/Explorer/Panes/ChangePartitionKeyPane/ChangePartitionKeyPane.tsx
index bf3f2d015..0bbc3722a 100644
--- a/src/Explorer/Panes/ChangePartitionKeyPane/ChangePartitionKeyPane.tsx
+++ b/src/Explorer/Panes/ChangePartitionKeyPane/ChangePartitionKeyPane.tsx
@@ -6,6 +6,7 @@ import {
   Icon,
   IconButton,
   Link,
+  MessageBar,
   Stack,
   Text,
   TooltipHost,
@@ -207,6 +208,7 @@ export const ChangePartitionKeyPane: React.FC<ChangePartitionKeyPaneProps> = ({
         </Stack>
         {createNewContainer ? (
           <Stack>
+            <MessageBar>All configurations except for unique keys will be copied from the source container</MessageBar>
             <Stack className="panelGroupSpacing">
               <Stack horizontal>
                 <span className="mandatoryStar">*&nbsp;</span>
diff --git a/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx b/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx
index 412b17051..6aaebc926 100644
--- a/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx
+++ b/src/Explorer/QueryCopilot/QueryCopilotPromptbar.tsx
@@ -214,8 +214,10 @@ export const QueryCopilotPromptbar: React.FC<QueryCopilotPromptProps> = ({
       const generateSQLQueryResponse: GenerateSQLQueryResponse = await response?.json();
       if (response.ok) {
         if (generateSQLQueryResponse?.sql !== "N/A") {
-          let currentGeneratedQuery = `-- **Prompt:** ${userPrompt}\r\n`;
-          currentGeneratedQuery += generateSQLQueryResponse.sql;
+          const queryExplanation = `-- **Explanation of query:** ${
+            generateSQLQueryResponse.explanation ? generateSQLQueryResponse.explanation : "N/A"
+          }\r\n`;
+          const currentGeneratedQuery = queryExplanation + generateSQLQueryResponse.sql;
           const lastQuery = generatedQuery && query ? `${query}\r\n` : "";
           setQuery(`${lastQuery}${currentGeneratedQuery}`);
           setGeneratedQuery(generateSQLQueryResponse.sql);