Partition Key Change with Container Copy (#1734)

* initial commit

* Add change partition key logic

* Update snapshot

* Update snapshot

* Update snapshot

* Update snapshot

* Cleanup code

* Disable Change on progress job

* add the database information in the panel

* add the database information in the panel

* clear in progress message and remove large partition key row

* hide from national cloud

* hide from national cloud

* Add check for public cloud

src/Common/dataAccess/dataTransfers.ts

@@ -0,0 +1,188 @@
+import { ApiType, userContext } from "UserContext";
+import * as NotificationConsoleUtils from "Utils/NotificationConsoleUtils";
+import {
+  cancel,
+  create,
+  get,
+  listByDatabaseAccount,
+} from "Utils/arm/generatedClients/dataTransferService/dataTransferJobs";
+import {
+  CosmosCassandraDataTransferDataSourceSink,
+  CosmosMongoDataTransferDataSourceSink,
+  CosmosSqlDataTransferDataSourceSink,
+  CreateJobRequest,
+  DataTransferJobFeedResults,
+  DataTransferJobGetResults,
+} from "Utils/arm/generatedClients/dataTransferService/types";
+import { addToPolling, removeFromPolling, updateDataTransferJob, useDataTransferJobs } from "hooks/useDataTransferJobs";
+import promiseRetry, { AbortError, FailedAttemptError } from "p-retry";
+
+export interface DataTransferParams {
+  jobName: string;
+  apiType: ApiType;
+  subscriptionId: string;
+  resourceGroupName: string;
+  accountName: string;
+  sourceDatabaseName: string;
+  sourceCollectionName: string;
+  targetDatabaseName: string;
+  targetCollectionName: string;
+}
+
+export const getDataTransferJobs = async (
+  subscriptionId: string,
+  resourceGroup: string,
+  accountName: string,
+): Promise<DataTransferJobGetResults[]> => {
+  let dataTransferJobs: DataTransferJobGetResults[] = [];
+  let dataTransferFeeds: DataTransferJobFeedResults = await listByDatabaseAccount(
+    subscriptionId,
+    resourceGroup,
+    accountName,
+  );
+  dataTransferJobs = [...dataTransferJobs, ...(dataTransferFeeds?.value || [])];
+  while (dataTransferFeeds?.nextLink) {
+    const nextResponse = await window.fetch(dataTransferFeeds.nextLink, {
+      headers: {
+        Authorization: userContext.authorizationToken,
+      },
+    });
+    if (nextResponse.ok) {
+      dataTransferFeeds = await nextResponse.json();
+      dataTransferJobs = [...dataTransferJobs, ...(dataTransferFeeds?.value || [])];
+    } else {
+      break;
+    }
+  }
+  return dataTransferJobs;
+};
+
+export const initiateDataTransfer = async (params: DataTransferParams): Promise<DataTransferJobGetResults> => {
+  const {
+    jobName,
+    apiType,
+    subscriptionId,
+    resourceGroupName,
+    accountName,
+    sourceDatabaseName,
+    sourceCollectionName,
+    targetDatabaseName,
+    targetCollectionName,
+  } = params;
+  const sourcePayload = createPayload(apiType, sourceDatabaseName, sourceCollectionName);
+  const targetPayload = createPayload(apiType, targetDatabaseName, targetCollectionName);
+  const body: CreateJobRequest = {
+    properties: {
+      source: sourcePayload,
+      destination: targetPayload,
+    },
+  };
+  return create(subscriptionId, resourceGroupName, accountName, jobName, body);
+};
+
+export const pollDataTransferJob = async (
+  jobName: string,
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+): Promise<unknown> => {
+  const currentPollingJobs = useDataTransferJobs.getState().pollingDataTransferJobs;
+  if (currentPollingJobs.has(jobName)) {
+    return;
+  }
+  let clearMessage = NotificationConsoleUtils.logConsoleProgress(`Data transfer job ${jobName} in progress`);
+  return await promiseRetry(
+    () => pollDataTransferJobOperation(jobName, subscriptionId, resourceGroupName, accountName, clearMessage),
+    {
+      retries: 500,
+      maxTimeout: 5000,
+      onFailedAttempt: (error: FailedAttemptError) => {
+        clearMessage();
+        clearMessage = NotificationConsoleUtils.logConsoleProgress(error.message);
+      },
+    },
+  );
+};
+
+const pollDataTransferJobOperation = async (
+  jobName: string,
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+  clearMessage?: () => void,
+): Promise<DataTransferJobGetResults> => {
+  if (!userContext.authorizationToken) {
+    throw new Error("No authority token provided");
+  }
+
+  addToPolling(jobName);
+
+  const body: DataTransferJobGetResults = await get(subscriptionId, resourceGroupName, accountName, jobName);
+  const status = body?.properties?.status;
+
+  updateDataTransferJob(body);
+
+  if (status === "Cancelled" || 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`);
+    throw new AbortError(error);
+  }
+  if (status === "Completed") {
+    removeFromPolling(jobName);
+    clearMessage && clearMessage();
+    NotificationConsoleUtils.logConsoleInfo(`Data transfer job ${jobName} completed`);
+    return body;
+  }
+  const processedCount = body.properties.processedCount;
+  const totalCount = body.properties.totalCount;
+  const retryMessage = `Data transfer job ${jobName} in progress, total count: ${totalCount}, processed count: ${processedCount}`;
+  throw new Error(retryMessage);
+};
+
+export const cancelDataTransferJob = async (
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+  jobName: string,
+): Promise<void> => {
+  const cancelResult: DataTransferJobGetResults = await cancel(subscriptionId, resourceGroupName, accountName, jobName);
+  updateDataTransferJob(cancelResult);
+  removeFromPolling(cancelResult?.properties?.jobName);
+};
+
+const createPayload = (
+  apiType: ApiType,
+  databaseName: string,
+  containerName: string,
+):
+  | CosmosSqlDataTransferDataSourceSink
+  | CosmosMongoDataTransferDataSourceSink
+  | CosmosCassandraDataTransferDataSourceSink => {
+  switch (apiType) {
+    case "SQL":
+      return {
+        component: "CosmosDBSql",
+        databaseName: databaseName,
+        containerName: containerName,
+      } as CosmosSqlDataTransferDataSourceSink;
+    case "Mongo":
+      return {
+        component: "CosmosDBMongo",
+        databaseName: databaseName,
+        collectionName: containerName,
+      } as CosmosMongoDataTransferDataSourceSink;
+    case "Cassandra":
+      return {
+        component: "CosmosDBCassandra",
+        keyspaceName: databaseName,
+        tableName: containerName,
+      };
+    default:
+      throw new Error(`Unsupported API type for data transfer: ${apiType}`);
+  }
+};

src/Explorer/Controls/Settings/SettingsComponent.tsx

@@ -1,5 +1,6 @@
 import { IPivotItemProps, IPivotProps, Pivot, PivotItem } from "@fluentui/react";
 import { useDatabases } from "Explorer/useDatabases";
+import { isRunningOnPublicCloud } from "Utils/CloudUtils";
 import * as React from "react";
 import DiscardIcon from "../../../../images/discard.svg";
 import SaveIcon from "../../../../images/save-cosmos.svg";
@@ -18,6 +19,10 @@ import { userContext } from "../../../UserContext";
 import * as AutoPilotUtils from "../../../Utils/AutoPilotUtils";
 import { MongoDBCollectionResource, MongoIndex } from "../../../Utils/arm/generatedClients/cosmos/types";
 import { CommandButtonComponentProps } from "../../Controls/CommandButton/CommandButtonComponent";
+import {
+  PartitionKeyComponent,
+  PartitionKeyComponentProps,
+} from "../../Controls/Settings/SettingsSubComponents/PartitionKeyComponent";
 import { useCommandBar } from "../../Menus/CommandBar/CommandBarComponentAdapter";
 import { SettingsTabV2 } from "../../Tabs/SettingsTabV2";
 import "./SettingsComponent.less";
@@ -128,6 +133,7 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
   private changeFeedPolicyVisible: boolean;
   private isFixedContainer: boolean;
   private shouldShowIndexingPolicyEditor: boolean;
+  private shouldShowPartitionKeyEditor: boolean;
   private totalThroughputUsed: number;
   public mongoDBCollectionResource: MongoDBCollectionResource;
 
@@ -140,6 +146,7 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       this.offer = this.collection?.offer();
       this.isAnalyticalStorageEnabled = !!this.collection?.analyticalStorageTtl();
       this.shouldShowIndexingPolicyEditor = userContext.apiType !== "Cassandra" && userContext.apiType !== "Mongo";
+      this.shouldShowPartitionKeyEditor = userContext.apiType === "SQL" && isRunningOnPublicCloud();
 
       this.changeFeedPolicyVisible = userContext.features.enableChangeFeedPolicy;
 
@@ -1056,6 +1063,12 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       onConflictResolutionDirtyChange: this.onConflictResolutionDirtyChange,
     };
 
+    const partitionKeyComponentProps: PartitionKeyComponentProps = {
+      database: useDatabases.getState().findDatabaseWithId(this.collection.databaseId),
+      collection: this.collection,
+      explorer: this.props.settingsTab.getContainer(),
+    };
+
     const tabs: SettingsV2TabInfo[] = [];
     if (!hasDatabaseSharedThroughput(this.collection) && this.offer) {
       tabs.push({
@@ -1091,6 +1104,13 @@ export class SettingsComponent extends React.Component<SettingsComponentProps, S
       });
     }
 
+    if (this.shouldShowPartitionKeyEditor) {
+      tabs.push({
+        tab: SettingsV2TabTypes.PartitionKeyTab,
+        content: <PartitionKeyComponent {...partitionKeyComponentProps} />,
+      });
+    }
+
     const pivotProps: IPivotProps = {
       onLinkClick: this.onPivotChange,
       selectedKey: SettingsV2TabTypes[this.state.selectedTab],

src/Explorer/Controls/Settings/SettingsSubComponents/PartitionKeyComponent.tsx

@@ -0,0 +1,216 @@
+import {
+  DefaultButton,
+  FontWeights,
+  Link,
+  MessageBar,
+  MessageBarType,
+  PrimaryButton,
+  ProgressIndicator,
+  Stack,
+  Text,
+} from "@fluentui/react";
+import * as React from "react";
+import * as ViewModels from "../../../../Contracts/ViewModels";
+
+import { handleError } from "Common/ErrorHandlingUtils";
+import { cancelDataTransferJob, pollDataTransferJob } from "Common/dataAccess/dataTransfers";
+import Explorer from "Explorer/Explorer";
+import { ChangePartitionKeyPane } from "Explorer/Panes/ChangePartitionKeyPane/ChangePartitionKeyPane";
+import {
+  CosmosSqlDataTransferDataSourceSink,
+  DataTransferJobGetResults,
+} from "Utils/arm/generatedClients/dataTransferService/types";
+import { refreshDataTransferJobs, useDataTransferJobs } from "hooks/useDataTransferJobs";
+import { useSidePanel } from "hooks/useSidePanel";
+import { userContext } from "../../../../UserContext";
+
+export interface PartitionKeyComponentProps {
+  database: ViewModels.Database;
+  collection: ViewModels.Collection;
+  explorer: Explorer;
+}
+
+export const PartitionKeyComponent: React.FC<PartitionKeyComponentProps> = ({ database, collection, explorer }) => {
+  const { dataTransferJobs } = useDataTransferJobs();
+  const [portalDataTransferJob, setPortalDataTransferJob] = React.useState<DataTransferJobGetResults>(null);
+
+  React.useEffect(() => {
+    const loadDataTransferJobs = refreshDataTransferOperations;
+    loadDataTransferJobs();
+  }, []);
+
+  React.useEffect(() => {
+    const currentJob = findPortalDataTransferJob();
+    setPortalDataTransferJob(currentJob);
+    startPollingforUpdate(currentJob);
+  }, [dataTransferJobs]);
+
+  const isHierarchicalPartitionedContainer = (): boolean => collection.partitionKey?.kind === "MultiHash";
+
+  const getPartitionKeyValue = (): string => {
+    return (collection.partitionKeyProperties || []).map((property) => "/" + property).join(", ");
+  };
+
+  const partitionKeyName = "Partition key";
+  const partitionKeyValue = getPartitionKeyValue();
+
+  const textHeadingStyle = {
+    root: { fontWeight: FontWeights.semibold, fontSize: 16 },
+  };
+
+  const textSubHeadingStyle = {
+    root: { fontWeight: FontWeights.semibold },
+  };
+
+  const startPollingforUpdate = (currentJob: DataTransferJobGetResults) => {
+    if (isCurrentJobInProgress(currentJob)) {
+      const jobName = currentJob?.properties?.jobName;
+      try {
+        pollDataTransferJob(
+          jobName,
+          userContext.subscriptionId,
+          userContext.resourceGroup,
+          userContext.databaseAccount.name,
+        );
+      } catch (error) {
+        handleError(error, "ChangePartitionKey", `Failed to complete data transfer job ${jobName}`);
+      }
+    }
+  };
+
+  const cancelRunningDataTransferJob = async (currentJob: DataTransferJobGetResults) => {
+    await cancelDataTransferJob(
+      userContext.subscriptionId,
+      userContext.resourceGroup,
+      userContext.databaseAccount.name,
+      currentJob?.properties?.jobName,
+    );
+  };
+
+  const isCurrentJobInProgress = (currentJob: DataTransferJobGetResults) => {
+    const jobStatus = currentJob?.properties?.status;
+    return (
+      jobStatus &&
+      jobStatus !== "Completed" &&
+      jobStatus !== "Cancelled" &&
+      jobStatus !== "Failed" &&
+      jobStatus !== "Faulted"
+    );
+  };
+
+  const refreshDataTransferOperations = async () => {
+    await refreshDataTransferJobs(
+      userContext.subscriptionId,
+      userContext.resourceGroup,
+      userContext.databaseAccount.name,
+    );
+  };
+
+  const findPortalDataTransferJob = (): DataTransferJobGetResults => {
+    return dataTransferJobs.find((feed: DataTransferJobGetResults) => {
+      const sourceSink: CosmosSqlDataTransferDataSourceSink = feed?.properties
+        ?.source as CosmosSqlDataTransferDataSourceSink;
+      return sourceSink.databaseName === collection.databaseId && sourceSink.containerName === collection.id();
+    });
+  };
+
+  const getProgressDescription = (): string => {
+    const processedCount = portalDataTransferJob?.properties?.processedCount;
+    const totalCount = portalDataTransferJob?.properties?.totalCount;
+    const processedCountString = totalCount > 0 ? `(${processedCount} of ${totalCount} documents processed)` : "";
+    return `${portalDataTransferJob?.properties?.status} ${processedCountString}`;
+  };
+
+  const startPartitionkeyChangeWorkflow = () => {
+    useSidePanel
+      .getState()
+      .openSidePanel(
+        "Change partition key",
+        <ChangePartitionKeyPane
+          sourceDatabase={database}
+          sourceCollection={collection}
+          explorer={explorer}
+          onClose={refreshDataTransferOperations}
+        />,
+      );
+  };
+
+  const getPercentageComplete = () => {
+    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;
+  };
+
+  return (
+    <Stack tokens={{ childrenGap: 20 }} styles={{ root: { maxWidth: 600 } }}>
+      <Stack tokens={{ childrenGap: 10 }}>
+        <Text styles={textHeadingStyle}>Change {partitionKeyName.toLowerCase()}</Text>
+        <Stack horizontal tokens={{ childrenGap: 20 }}>
+          <Stack tokens={{ childrenGap: 5 }}>
+            <Text styles={textSubHeadingStyle}>Current {partitionKeyName.toLowerCase()}</Text>
+            <Text styles={textSubHeadingStyle}>Partitioning</Text>
+          </Stack>
+          <Stack tokens={{ childrenGap: 5 }}>
+            <Text>{partitionKeyValue}</Text>
+            <Text>{isHierarchicalPartitionedContainer() ? "Hierarchical" : "Non-hierarchical"}</Text>
+          </Stack>
+        </Stack>
+      </Stack>
+      <MessageBar messageBarType={MessageBarType.warning}>
+        To safeguard the integrity of the data being copied to the new container, ensure that no updates are made to the
+        source container for the entire duration of the partition key change process.
+        <Link
+          href="https://learn.microsoft.com/azure/cosmos-db/container-copy#how-does-container-copy-work"
+          target="_blank"
+          underline
+        >
+          Learn more
+        </Link>
+      </MessageBar>
+      <Text>
+        To change the partition key, a new destination container must be created or an existing destination container
+        selected. Data will then be copied to the destination container.
+      </Text>
+      <PrimaryButton
+        styles={{ root: { width: "fit-content" } }}
+        text="Change"
+        onClick={startPartitionkeyChangeWorkflow}
+        disabled={isCurrentJobInProgress(portalDataTransferJob)}
+      />
+      {portalDataTransferJob && (
+        <Stack>
+          <Text styles={textHeadingStyle}>{partitionKeyName} change job</Text>
+          <Stack
+            horizontal
+            tokens={{ childrenGap: 20 }}
+            styles={{
+              root: {
+                alignItems: "center",
+              },
+            }}
+          >
+            <ProgressIndicator
+              label={portalDataTransferJob?.properties?.jobName}
+              description={getProgressDescription()}
+              percentComplete={getPercentageComplete()}
+              styles={{
+                root: {
+                  width: "85%",
+                },
+              }}
+            ></ProgressIndicator>
+            {isCurrentJobInProgress(portalDataTransferJob) && (
+              <DefaultButton text="Cancel" onClick={() => cancelRunningDataTransferJob(portalDataTransferJob)} />
+            )}
+          </Stack>
+        </Stack>
+      )}
+    </Stack>
+  );
+};

src/Explorer/Controls/Settings/SettingsUtils.tsx

@@ -45,6 +45,7 @@ export enum SettingsV2TabTypes {
   ConflictResolutionTab,
   SubSettingsTab,
   IndexingPolicyTab,
+  PartitionKeyTab,
 }
 
 export interface IsComponentDirtyResult {
@@ -146,6 +147,8 @@ export const getTabTitle = (tab: SettingsV2TabTypes): string => {
       return "Settings";
     case SettingsV2TabTypes.IndexingPolicyTab:
       return "Indexing Policy";
+    case SettingsV2TabTypes.PartitionKeyTab:
+      return "Partition Keys";
     default:
       throw new Error(`Unknown tab ${tab}`);
   }
@@ -199,3 +202,49 @@ export const getMongoIndexTypeText = (index: MongoIndexTypes): string => {
 export const isIndexTransforming = (indexTransformationProgress: number): boolean =>
   // index transformation progress can be 0
   indexTransformationProgress !== undefined && indexTransformationProgress !== 100;
+
+export const getPartitionKeyName = (apiType: string, isLowerCase?: boolean): string => {
+  const partitionKeyName = apiType === "Mongo" ? "Shard key" : "Partition key";
+  return isLowerCase ? partitionKeyName.toLocaleLowerCase() : partitionKeyName;
+};
+
+export const getPartitionKeyTooltipText = (apiType: string): string => {
+  if (apiType === "Mongo") {
+    return "The shard key (field) is used to split your data across many replica sets (shards) to achieve unlimited scalability. It’s critical to choose a field that will evenly distribute your data.";
+  }
+  let tooltipText = `The ${getPartitionKeyName(
+    apiType,
+    true,
+  )} is used to automatically distribute data across partitions for scalability. Choose a property in your JSON document that has a wide range of values and evenly distributes request volume.`;
+  if (apiType === "SQL") {
+    tooltipText += " For small read-heavy workloads or write-heavy workloads of any size, id is often a good choice.";
+  }
+  return tooltipText;
+};
+
+export const getPartitionKeySubtext = (partitionKeyDefault: boolean, apiType: string): string => {
+  if (partitionKeyDefault && (apiType === "SQL" || apiType === "Mongo")) {
+    const subtext = "For small workloads, the item ID is a suitable choice for the partition key.";
+    return subtext;
+  }
+  return "";
+};
+
+export const getPartitionKeyPlaceHolder = (apiType: string, index?: number): string => {
+  switch (apiType) {
+    case "Mongo":
+      return "e.g., categoryId";
+    case "Gremlin":
+      return "e.g., /address";
+    case "SQL":
+      return `${
+        index === undefined
+          ? "Required - first partition key e.g., /TenantId"
+          : index === 0
+          ? "second partition key e.g., /UserId"
+          : "third partition key e.g., /SessionId"
+      }`;
+    default:
+      return "e.g., /address/zipCode";
+  }
+};

src/Explorer/Controls/Settings/__snapshots__/SettingsComponent.test.tsx.snap

@@ -204,6 +204,98 @@ exports[`SettingsComponent renders 1`] = `
           shouldDiscardIndexingPolicy={false}
         />
       </PivotItem>
+      <PivotItem
+        headerText="Partition Keys"
+        itemKey="PartitionKeyTab"
+        key="PartitionKeyTab"
+        style={
+          Object {
+            "marginTop": 20,
+          }
+        }
+      >
+        <PartitionKeyComponent
+          collection={
+            Object {
+              "analyticalStorageTtl": [Function],
+              "changeFeedPolicy": [Function],
+              "conflictResolutionPolicy": [Function],
+              "container": Explorer {
+                "_isInitializingNotebooks": false,
+                "_resetNotebookWorkspace": [Function],
+                "isFixedCollectionWithSharedThroughputSupported": [Function],
+                "isTabsContentExpanded": [Function],
+                "onRefreshDatabasesKeyPress": [Function],
+                "onRefreshResourcesClick": [Function],
+                "phoenixClient": PhoenixClient {
+                  "armResourceId": undefined,
+                  "retryOptions": Object {
+                    "maxTimeout": 5000,
+                    "minTimeout": 5000,
+                    "retries": 3,
+                  },
+                },
+                "provideFeedbackEmail": [Function],
+                "queriesClient": QueriesClient {
+                  "container": [Circular],
+                },
+                "refreshNotebookList": [Function],
+                "resourceTree": ResourceTreeAdapter {
+                  "container": [Circular],
+                  "copyNotebook": [Function],
+                  "parameters": [Function],
+                },
+              },
+              "databaseId": "test",
+              "defaultTtl": [Function],
+              "geospatialConfig": [Function],
+              "getDatabase": [Function],
+              "id": [Function],
+              "indexingPolicy": [Function],
+              "offer": [Function],
+              "partitionKey": Object {
+                "kind": "hash",
+                "paths": Array [],
+                "version": 2,
+              },
+              "partitionKeyProperties": Array [
+                "partitionKey",
+              ],
+              "readSettings": [Function],
+              "uniqueKeyPolicy": Object {},
+              "usageSizeInKB": [Function],
+            }
+          }
+          explorer={
+            Explorer {
+              "_isInitializingNotebooks": false,
+              "_resetNotebookWorkspace": [Function],
+              "isFixedCollectionWithSharedThroughputSupported": [Function],
+              "isTabsContentExpanded": [Function],
+              "onRefreshDatabasesKeyPress": [Function],
+              "onRefreshResourcesClick": [Function],
+              "phoenixClient": PhoenixClient {
+                "armResourceId": undefined,
+                "retryOptions": Object {
+                  "maxTimeout": 5000,
+                  "minTimeout": 5000,
+                  "retries": 3,
+                },
+              },
+              "provideFeedbackEmail": [Function],
+              "queriesClient": QueriesClient {
+                "container": [Circular],
+              },
+              "refreshNotebookList": [Function],
+              "resourceTree": ResourceTreeAdapter {
+                "container": [Circular],
+                "copyNotebook": [Function],
+                "parameters": [Function],
+              },
+            }
+          }
+        />
+      </PivotItem>
     </StyledPivot>
   </div>
 </div>

src/Explorer/Panes/ChangePartitionKeyPane/ChangePartitionKeyPane.tsx

@@ -0,0 +1,396 @@
+import {
+  DefaultButton,
+  DirectionalHint,
+  Dropdown,
+  IDropdownOption,
+  Icon,
+  IconButton,
+  Link,
+  Stack,
+  Text,
+  TooltipHost,
+} from "@fluentui/react";
+import * as Constants from "Common/Constants";
+import { handleError } from "Common/ErrorHandlingUtils";
+import { createCollection } from "Common/dataAccess/createCollection";
+import { DataTransferParams, initiateDataTransfer } from "Common/dataAccess/dataTransfers";
+import * as DataModels from "Contracts/DataModels";
+import * as ViewModels from "Contracts/ViewModels";
+import {
+  getPartitionKeyName,
+  getPartitionKeyPlaceHolder,
+  getPartitionKeySubtext,
+  getPartitionKeyTooltipText,
+} from "Explorer/Controls/Settings/SettingsUtils";
+import Explorer from "Explorer/Explorer";
+import { RightPaneForm } from "Explorer/Panes/RightPaneForm/RightPaneForm";
+import { useDatabases } from "Explorer/useDatabases";
+import { userContext } from "UserContext";
+import { getCollectionName } from "Utils/APITypeUtils";
+import { useSidePanel } from "hooks/useSidePanel";
+import * as React from "react";
+
+export interface ChangePartitionKeyPaneProps {
+  sourceDatabase: ViewModels.Database;
+  sourceCollection: ViewModels.Collection;
+  explorer: Explorer;
+  onClose: () => Promise<void>;
+}
+
+export const ChangePartitionKeyPane: React.FC<ChangePartitionKeyPaneProps> = ({
+  sourceDatabase,
+  sourceCollection,
+  explorer,
+  onClose,
+}) => {
+  const [targetCollectionId, setTargetCollectionId] = React.useState<string>();
+  const [createNewContainer, setCreateNewContainer] = React.useState<boolean>(true);
+  const [formError, setFormError] = React.useState<string>();
+  const [isExecuting, setIsExecuting] = React.useState<boolean>(false);
+  const [subPartitionKeys, setSubPartitionKeys] = React.useState<string[]>([]);
+  const [partitionKey, setPartitionKey] = React.useState<string>();
+
+  const getCollectionOptions = (): IDropdownOption[] => {
+    return sourceDatabase
+      .collections()
+      .filter((collection) => collection.id !== sourceCollection.id)
+      .map((collection) => ({
+        key: collection.id(),
+        text: collection.id(),
+      }));
+  };
+
+  const submit = async () => {
+    if (!validateInputs()) {
+      return;
+    }
+    setIsExecuting(true);
+    try {
+      createNewContainer && (await createContainer());
+      await createDataTransferJob();
+      await onClose();
+    } catch (error) {
+      handleError(error, "ChangePartitionKey", "Failed to start data transfer job");
+    }
+    setIsExecuting(false);
+    useSidePanel.getState().closeSidePanel();
+  };
+
+  const validateInputs = (): boolean => {
+    if (!createNewContainer && !targetCollectionId) {
+      setFormError("Choose an existing container");
+      return false;
+    }
+    return true;
+  };
+
+  const createDataTransferJob = async () => {
+    const jobName = `Portal_${targetCollectionId}_${Math.floor(Date.now() / 1000)}`;
+    const dataTransferParams: DataTransferParams = {
+      jobName,
+      apiType: userContext.apiType,
+      subscriptionId: userContext.subscriptionId,
+      resourceGroupName: userContext.resourceGroup,
+      accountName: userContext.databaseAccount.name,
+      sourceDatabaseName: sourceDatabase.id(),
+      sourceCollectionName: sourceCollection.id(),
+      targetDatabaseName: sourceDatabase.id(),
+      targetCollectionName: targetCollectionId,
+    };
+    await initiateDataTransfer(dataTransferParams);
+  };
+
+  const createContainer = async () => {
+    const partitionKeyString = partitionKey.trim();
+    const partitionKeyData: DataModels.PartitionKey = partitionKeyString
+      ? {
+          paths: [partitionKeyString, ...(subPartitionKeys.length > 0 ? subPartitionKeys : [])],
+          kind: subPartitionKeys.length > 0 ? "MultiHash" : "Hash",
+          version: 2,
+        }
+      : undefined;
+
+    const createCollectionParams: DataModels.CreateCollectionParams = {
+      createNewDatabase: false,
+      collectionId: targetCollectionId,
+      databaseId: sourceDatabase.id(),
+      databaseLevelThroughput: isSelectedDatabaseSharedThroughput(),
+      offerThroughput: sourceCollection.offer()?.manualThroughput,
+      autoPilotMaxThroughput: sourceCollection.offer()?.autoscaleMaxThroughput,
+      partitionKey: partitionKeyData,
+    };
+    await createCollection(createCollectionParams);
+    await explorer.refreshAllDatabases();
+  };
+
+  const isSelectedDatabaseSharedThroughput = (): boolean => {
+    const selectedDatabase = useDatabases
+      .getState()
+      .databases?.find((database) => database.id() === sourceDatabase.id());
+    return !!selectedDatabase?.offer();
+  };
+
+  return (
+    <RightPaneForm formError={formError} isExecuting={isExecuting} onSubmit={submit} submitButtonText="OK">
+      <Stack tokens={{ childrenGap: 10 }} className="panelMainContent">
+        <Text variant="small">
+          When changing a container’s partition key, you will need to create a destination container with the correct
+          partition key. You may also select an existing destination container.&nbsp;
+          <Link
+            href="https://learn.microsoft.com/en-us/azure/cosmos-db/container-copy#container-copy-within-an-azure-cosmos-db-account"
+            target="_blank"
+            underline
+          >
+            Learn more
+          </Link>
+        </Text>
+        <Stack>
+          <Stack horizontal>
+            <span className="mandatoryStar">*&nbsp;</span>
+            <Text className="panelTextBold" variant="small">
+              Database id
+            </Text>
+            <TooltipHost
+              directionalHint={DirectionalHint.bottomLeftEdge}
+              content={`A database is analogous to a namespace. It is the unit of management for a set of ${getCollectionName(
+                true,
+              ).toLocaleLowerCase()}.`}
+            >
+              <Icon
+                iconName="Info"
+                className="panelInfoIcon"
+                tabIndex={0}
+                aria-label={`A database is analogous to a namespace. It is the unit of management for a set of ${getCollectionName(
+                  true,
+                ).toLocaleLowerCase()}.`}
+              />
+            </TooltipHost>
+          </Stack>
+          <Dropdown
+            styles={{ title: { height: 27, lineHeight: 27 }, dropdownItem: { fontSize: 12 } }}
+            style={{ width: 300, fontSize: 12 }}
+            options={[]}
+            placeholder={sourceDatabase.id()}
+            responsiveMode={999}
+            disabled={true}
+          />
+        </Stack>
+        <Stack className="panelGroupSpacing" horizontal verticalAlign="center">
+          <div role="radiogroup">
+            <input
+              className="panelRadioBtn"
+              checked={createNewContainer}
+              aria-label="Create new container"
+              aria-checked={createNewContainer}
+              name="containerType"
+              type="radio"
+              role="radio"
+              id="containerCreateNew"
+              tabIndex={0}
+              onChange={() => setCreateNewContainer(true)}
+            />
+            <span className="panelRadioBtnLabel">New container</span>
+
+            <input
+              className="panelRadioBtn"
+              checked={!createNewContainer}
+              aria-label="Use existing container"
+              aria-checked={!createNewContainer}
+              name="containerType"
+              type="radio"
+              role="radio"
+              tabIndex={0}
+              onChange={() => setCreateNewContainer(false)}
+            />
+            <span className="panelRadioBtnLabel">Existing container</span>
+          </div>
+        </Stack>
+        {createNewContainer ? (
+          <Stack>
+            <Stack className="panelGroupSpacing">
+              <Stack horizontal>
+                <span className="mandatoryStar">*&nbsp;</span>
+                <Text className="panelTextBold" variant="small">
+                  {`${getCollectionName()} id`}
+                </Text>
+                <TooltipHost
+                  directionalHint={DirectionalHint.bottomLeftEdge}
+                  content={`Unique identifier for the ${getCollectionName().toLocaleLowerCase()} and used for id-based routing through REST and all SDKs.`}
+                >
+                  <Icon
+                    role="button"
+                    iconName="Info"
+                    className="panelInfoIcon"
+                    tabIndex={0}
+                    ariaLabel={`Unique identifier for the ${getCollectionName().toLocaleLowerCase()} and used for id-based routing through REST and all SDKs.`}
+                  />
+                </TooltipHost>
+              </Stack>
+              <input
+                name="collectionId"
+                id="collectionId"
+                type="text"
+                aria-required
+                required
+                autoComplete="off"
+                pattern="[^/?#\\]*[^/?# \\]"
+                title="May not end with space nor contain characters '\' '/' '#' '?'"
+                placeholder={`e.g., ${getCollectionName()}1`}
+                size={40}
+                className="panelTextField"
+                aria-label={`${getCollectionName()} id, Example ${getCollectionName()}1`}
+                value={targetCollectionId}
+                onChange={(event: React.ChangeEvent<HTMLInputElement>) => setTargetCollectionId(event.target.value)}
+              />
+            </Stack>
+            <Stack tokens={{ childrenGap: 10 }}>
+              <Stack horizontal>
+                <span className="mandatoryStar">*&nbsp;</span>
+                <Text className="panelTextBold" variant="small">
+                  {getPartitionKeyName(userContext.apiType)}
+                </Text>
+                <TooltipHost
+                  directionalHint={DirectionalHint.bottomLeftEdge}
+                  content={getPartitionKeyTooltipText(userContext.apiType)}
+                >
+                  <Icon
+                    iconName="Info"
+                    className="panelInfoIcon"
+                    tabIndex={0}
+                    aria-label={getPartitionKeyTooltipText(userContext.apiType)}
+                  />
+                </TooltipHost>
+              </Stack>
+
+              <Text variant="small" aria-label="pkDescription">
+                {getPartitionKeySubtext(userContext.features.partitionKeyDefault, userContext.apiType)}
+              </Text>
+
+              <input
+                type="text"
+                id="addCollection-partitionKeyValue"
+                aria-required
+                required
+                size={40}
+                className="panelTextField"
+                placeholder={getPartitionKeyPlaceHolder(userContext.apiType)}
+                aria-label={getPartitionKeyName(userContext.apiType)}
+                pattern={".*"}
+                title={""}
+                value={partitionKey}
+                onChange={(event: React.ChangeEvent<HTMLInputElement>) => {
+                  if (!partitionKey && !event.target.value.startsWith("/")) {
+                    setPartitionKey("/" + event.target.value);
+                  } else {
+                    setPartitionKey(event.target.value);
+                  }
+                }}
+              />
+              {subPartitionKeys.map((subPartitionKey: string, index: number) => {
+                return (
+                  <Stack style={{ marginBottom: 8 }} key={`uniqueKey${index}`} horizontal>
+                    <div
+                      style={{
+                        width: "20px",
+                        border: "solid",
+                        borderWidth: "0px 0px 1px 1px",
+                        marginRight: "5px",
+                      }}
+                    ></div>
+                    <input
+                      type="text"
+                      id="addCollection-partitionKeyValue"
+                      key={`addCollection-partitionKeyValue_${index}`}
+                      aria-required
+                      required
+                      size={40}
+                      tabIndex={index > 0 ? 1 : 0}
+                      className="panelTextField"
+                      autoComplete="off"
+                      placeholder={getPartitionKeyPlaceHolder(userContext.apiType, index)}
+                      aria-label={getPartitionKeyName(userContext.apiType)}
+                      pattern={".*"}
+                      title={""}
+                      value={subPartitionKey}
+                      onChange={(event: React.ChangeEvent<HTMLInputElement>) => {
+                        const keys = [...subPartitionKeys];
+                        if (!keys[index] && !event.target.value.startsWith("/")) {
+                          keys[index] = "/" + event.target.value.trim();
+                          setSubPartitionKeys(keys);
+                        } else {
+                          keys[index] = event.target.value.trim();
+                          setSubPartitionKeys(keys);
+                        }
+                      }}
+                    />
+                    <IconButton
+                      iconProps={{ iconName: "Delete" }}
+                      style={{ height: 27 }}
+                      onClick={() => {
+                        const keys = subPartitionKeys.filter((uniqueKey, j) => index !== j);
+                        setSubPartitionKeys(keys);
+                      }}
+                    />
+                  </Stack>
+                );
+              })}
+              <Stack className="panelGroupSpacing">
+                <DefaultButton
+                  styles={{ root: { padding: 0, width: 200, height: 30 }, label: { fontSize: 12 } }}
+                  disabled={subPartitionKeys.length >= Constants.BackendDefaults.maxNumMultiHashPartition}
+                  onClick={() => setSubPartitionKeys([...subPartitionKeys, ""])}
+                >
+                  Add hierarchical partition key
+                </DefaultButton>
+                {subPartitionKeys.length > 0 && (
+                  <Text variant="small">
+                    <Icon iconName="InfoSolid" className="removeIcon" tabIndex={0} /> This feature allows you to
+                    partition your data with up to three levels of keys for better data distribution. Requires .NET V3,
+                    Java V4 SDK, or preview JavaScript V3 SDK.{" "}
+                    <Link href="https://aka.ms/cosmos-hierarchical-partitioning" target="_blank">
+                      Learn more
+                    </Link>
+                  </Text>
+                )}
+              </Stack>
+            </Stack>
+          </Stack>
+        ) : (
+          <Stack>
+            <Stack horizontal>
+              <span className="mandatoryStar">*&nbsp;</span>
+              <Text className="panelTextBold" variant="small">
+                {`${getCollectionName()}`}
+              </Text>
+              <TooltipHost
+                directionalHint={DirectionalHint.bottomLeftEdge}
+                content={`Unique identifier for the ${getCollectionName().toLocaleLowerCase()} and used for id-based routing through REST and all SDKs.`}
+              >
+                <Icon
+                  role="button"
+                  iconName="Info"
+                  className="panelInfoIcon"
+                  tabIndex={0}
+                  ariaLabel={`Unique identifier for the ${getCollectionName().toLocaleLowerCase()} and used for id-based routing through REST and all SDKs.`}
+                />
+              </TooltipHost>
+            </Stack>
+
+            <Dropdown
+              styles={{ title: { height: 27, lineHeight: 27 }, dropdownItem: { fontSize: 12 } }}
+              style={{ width: 300, fontSize: 12 }}
+              placeholder="Choose an existing container"
+              options={getCollectionOptions()}
+              onChange={(event: React.FormEvent<HTMLDivElement>, collection: IDropdownOption) => {
+                setTargetCollectionId(collection.key as string);
+                setFormError("");
+              }}
+              defaultSelectedKey={targetCollectionId}
+              responsiveMode={999}
+            />
+          </Stack>
+        )}
+      </Stack>
+    </RightPaneForm>
+  );
+};

src/UserContext.ts

@@ -86,7 +86,7 @@ interface UserContext {
 }
 
 export type ApiType = "SQL" | "Mongo" | "Gremlin" | "Tables" | "Cassandra" | "Postgres" | "VCoreMongo";
-export type PortalEnv = "localhost" | "blackforest" | "fairfax" | "mooncake" | "prod" | "dev";
+export type PortalEnv = "localhost" | "blackforest" | "fairfax" | "mooncake" | "prod1" | "rx" | "ex" | "prod" | "dev";
 
 const ONE_WEEK_IN_MS = 604800000;
 

src/Utils/CloudUtils.ts

@@ -1,9 +1,9 @@
 import { userContext } from "../UserContext";
 
 export function isRunningOnNationalCloud(): boolean {
-  return (
+  return !isRunningOnPublicCloud();
-    userContext.portalEnv === "blackforest" ||
+}
-    userContext.portalEnv === "fairfax" ||
+
-    userContext.portalEnv === "mooncake"
+export function isRunningOnPublicCloud(): boolean {
-  );
+  return userContext?.portalEnv === "prod1" || userContext?.portalEnv === "prod";
 }

src/Utils/arm/generatedClients/dataTransferService/dataTransferJobs.ts

@@ -0,0 +1,78 @@
+/* 
+  AUTOGENERATED FILE
+  Run "npm run generateARMClients" to regenerate
+  Edting this file directly should be done with extreme caution as not to diverge from ARM REST specs
+
+  Generated from: https://raw.githubusercontent.com/Azure/azure-rest-api-specs/main/specification/cosmos-db/resource-manager/Microsoft.DocumentDB/preview/2023-11-15-preview/dataTransferService.json
+*/
+
+import { configContext } from "../../../../ConfigContext";
+import { armRequest } from "../../request";
+import * as Types from "./types";
+const apiVersion = "2023-11-15-preview";
+
+/* Creates a Data Transfer Job. */
+export async function create(
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+  jobName: string,
+  body: Types.CreateJobRequest,
+): Promise<Types.DataTransferJobGetResults> {
+  const path = `/subscriptions/${subscriptionId}/resourceGroups/${resourceGroupName}/providers/Microsoft.DocumentDB/databaseAccounts/${accountName}/dataTransferJobs/${jobName}`;
+  return armRequest({ host: configContext.ARM_ENDPOINT, path, method: "PUT", apiVersion, body });
+}
+
+/* Get a Data Transfer Job. */
+export async function get(
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+  jobName: string,
+): Promise<Types.DataTransferJobGetResults> {
+  const path = `/subscriptions/${subscriptionId}/resourceGroups/${resourceGroupName}/providers/Microsoft.DocumentDB/databaseAccounts/${accountName}/dataTransferJobs/${jobName}`;
+  return armRequest({ host: configContext.ARM_ENDPOINT, path, method: "GET", apiVersion });
+}
+
+/* Pause a Data Transfer Job. */
+export async function pause(
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+  jobName: string,
+): Promise<Types.DataTransferJobGetResults> {
+  const path = `/subscriptions/${subscriptionId}/resourceGroups/${resourceGroupName}/providers/Microsoft.DocumentDB/databaseAccounts/${accountName}/dataTransferJobs/${jobName}/pause`;
+  return armRequest({ host: configContext.ARM_ENDPOINT, path, method: "POST", apiVersion });
+}
+
+/* Resumes a Data Transfer Job. */
+export async function resume(
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+  jobName: string,
+): Promise<Types.DataTransferJobGetResults> {
+  const path = `/subscriptions/${subscriptionId}/resourceGroups/${resourceGroupName}/providers/Microsoft.DocumentDB/databaseAccounts/${accountName}/dataTransferJobs/${jobName}/resume`;
+  return armRequest({ host: configContext.ARM_ENDPOINT, path, method: "POST", apiVersion });
+}
+
+/* Cancels a Data Transfer Job. */
+export async function cancel(
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+  jobName: string,
+): Promise<Types.DataTransferJobGetResults> {
+  const path = `/subscriptions/${subscriptionId}/resourceGroups/${resourceGroupName}/providers/Microsoft.DocumentDB/databaseAccounts/${accountName}/dataTransferJobs/${jobName}/cancel`;
+  return armRequest({ host: configContext.ARM_ENDPOINT, path, method: "POST", apiVersion });
+}
+
+/* Get a list of Data Transfer jobs. */
+export async function listByDatabaseAccount(
+  subscriptionId: string,
+  resourceGroupName: string,
+  accountName: string,
+): Promise<Types.DataTransferJobFeedResults> {
+  const path = `/subscriptions/${subscriptionId}/resourceGroups/${resourceGroupName}/providers/Microsoft.DocumentDB/databaseAccounts/${accountName}/dataTransferJobs`;
+  return armRequest({ host: configContext.ARM_ENDPOINT, path, method: "GET", apiVersion });
+}

src/Utils/arm/generatedClients/dataTransferService/types.ts

@@ -0,0 +1,101 @@
+/* 
+  AUTOGENERATED FILE
+  Run "npm run generateARMClients" to regenerate
+  Edting this file directly should be done with extreme caution as not to diverge from ARM REST specs
+
+  Generated from: https://raw.githubusercontent.com/Azure/azure-rest-api-specs/main/specification/cosmos-db/resource-manager/Microsoft.DocumentDB/preview/2023-11-15-preview/dataTransferService.json
+*/
+
+/* Base class for all DataTransfer source/sink */
+export interface DataTransferDataSourceSink {
+  /* undocumented */
+  component: "CosmosDBCassandra" | "CosmosDBMongo" | "CosmosDBSql" | "AzureBlobStorage";
+}
+
+/* A base CosmosDB data source/sink */
+export type BaseCosmosDataTransferDataSourceSink = DataTransferDataSourceSink & {
+  /* undocumented */
+  remoteAccountName?: string;
+};
+
+/* A CosmosDB Cassandra API data source/sink */
+export type CosmosCassandraDataTransferDataSourceSink = BaseCosmosDataTransferDataSourceSink & {
+  /* undocumented */
+  keyspaceName: string;
+  /* undocumented */
+  tableName: string;
+};
+
+/* A CosmosDB Mongo API data source/sink */
+export type CosmosMongoDataTransferDataSourceSink = BaseCosmosDataTransferDataSourceSink & {
+  /* undocumented */
+  databaseName: string;
+  /* undocumented */
+  collectionName: string;
+};
+
+/* A CosmosDB No Sql API data source/sink */
+export type CosmosSqlDataTransferDataSourceSink = BaseCosmosDataTransferDataSourceSink & {
+  /* undocumented */
+  databaseName: string;
+  /* undocumented */
+  containerName: string;
+};
+
+/* An Azure Blob Storage data source/sink */
+export type AzureBlobDataTransferDataSourceSink = DataTransferDataSourceSink & {
+  /* undocumented */
+  containerName: string;
+  /* undocumented */
+  endpointUrl?: string;
+};
+
+/* The properties of a DataTransfer Job */
+export interface DataTransferJobProperties {
+  /* Job Name */
+  readonly jobName?: string;
+  /* Source DataStore details */
+  source: DataTransferDataSourceSink;
+
+  /* Destination DataStore details */
+  destination: DataTransferDataSourceSink;
+
+  /* Job Status */
+  readonly status?: string;
+  /* Processed Count. */
+  readonly processedCount?: number;
+  /* Total Count. */
+  readonly totalCount?: number;
+  /* Last Updated Time (ISO-8601 format). */
+  readonly lastUpdatedUtcTime?: string;
+  /* Worker count */
+  workerCount?: number;
+  /* Error response for Faulted job */
+  readonly error?: unknown;
+
+  /* Total Duration of Job */
+  readonly duration?: string;
+  /* Mode of job execution */
+  mode?: "Offline" | "Online";
+}
+
+/* Parameters to create Data Transfer Job */
+export type CreateJobRequest = unknown & {
+  /* Data Transfer Create Job Properties */
+  properties: DataTransferJobProperties;
+};
+
+/* A Cosmos DB Data Transfer Job */
+export type DataTransferJobGetResults = unknown & {
+  /* undocumented */
+  properties?: DataTransferJobProperties;
+};
+
+/* The List operation response, that contains the Data Transfer jobs and their properties. */
+export interface DataTransferJobFeedResults {
+  /* List of Data Transfer jobs and their properties. */
+  readonly value?: DataTransferJobGetResults[];
+
+  /* URL to get the next set of Data Transfer job list results if there are any. */
+  readonly nextLink?: string;
+}

src/hooks/useDataTransferJobs.tsx

@@ -0,0 +1,60 @@
+import { getDataTransferJobs } from "Common/dataAccess/dataTransfers";
+import { DataTransferJobGetResults } from "Utils/arm/generatedClients/dataTransferService/types";
+import create, { UseStore } from "zustand";
+
+export interface DataTransferJobsState {
+  dataTransferJobs: DataTransferJobGetResults[];
+  pollingDataTransferJobs: Set<string>;
+  setDataTransferJobs: (dataTransferJobs: DataTransferJobGetResults[]) => void;
+  setPollingDataTransferJobs: (pollingDataTransferJobs: Set<string>) => void;
+}
+
+type DataTransferJobStore = UseStore<DataTransferJobsState>;
+
+export const useDataTransferJobs: DataTransferJobStore = create((set) => ({
+  dataTransferJobs: [],
+  pollingDataTransferJobs: new Set<string>(),
+  setDataTransferJobs: (dataTransferJobs: DataTransferJobGetResults[]) => set({ dataTransferJobs }),
+  setPollingDataTransferJobs: (pollingDataTransferJobs: Set<string>) => set({ pollingDataTransferJobs }),
+}));
+
+export const refreshDataTransferJobs = async (
+  subscriptionId: string,
+  resourceGroup: string,
+  accountName: string,
+): Promise<DataTransferJobGetResults[]> => {
+  const dataTransferJobs: DataTransferJobGetResults[] = await getDataTransferJobs(
+    subscriptionId,
+    resourceGroup,
+    accountName,
+  );
+  const jobRegex = /^Portal_(.+)_(\d{10,})$/;
+  const sortedJobs: DataTransferJobGetResults[] = dataTransferJobs?.sort(
+    (a, b) =>
+      new Date(b?.properties?.lastUpdatedUtcTime).getTime() - new Date(a?.properties?.lastUpdatedUtcTime).getTime(),
+  );
+  const filteredJobs = sortedJobs.filter((job) => jobRegex.test(job?.properties?.jobName));
+  useDataTransferJobs.getState().setDataTransferJobs(filteredJobs);
+  return filteredJobs;
+};
+
+export const updateDataTransferJob = (updateJob: DataTransferJobGetResults) => {
+  const updatedDataTransferJobs = useDataTransferJobs
+    .getState()
+    .dataTransferJobs.map((job: DataTransferJobGetResults) =>
+      job?.properties?.jobName === updateJob?.properties?.jobName ? updateJob : job,
+    );
+  useDataTransferJobs.getState().setDataTransferJobs(updatedDataTransferJobs);
+};
+
+export const addToPolling = (addJob: string) => {
+  const pollingJobs = useDataTransferJobs.getState().pollingDataTransferJobs;
+  pollingJobs.add(addJob);
+  useDataTransferJobs.getState().setPollingDataTransferJobs(pollingJobs);
+};
+
+export const removeFromPolling = (removeJob: string) => {
+  const pollingJobs = useDataTransferJobs.getState().pollingDataTransferJobs;
+  pollingJobs.delete(removeJob);
+  useDataTransferJobs.getState().setPollingDataTransferJobs(pollingJobs);
+};

utils/armClientGenerator/generator.ts

@@ -16,13 +16,13 @@ Results of this file should be checked into the repo.
 */
 
 // CHANGE THESE VALUES TO GENERATE NEW CLIENTS
-const version = "2023-09-15-preview";
+const version = "2023-11-15-preview";
 /* The following are legal options for resourceName but you generally will only use cosmos-db:
 "cosmos-db" | "managedCassandra" | "mongorbac" | "notebook" | "privateEndpointConnection" | "privateLinkResources" |
-"rbac" | "restorable" | "services"
+"rbac" | "restorable" | "services" | "dataTransferService"
 */
 const githubResourceName = "cosmos-db";
-const deResourceName = "cosmos";
+const deResourceName = "cosmos-db";
 const schemaURL = `https://raw.githubusercontent.com/Azure/azure-rest-api-specs/main/specification/cosmos-db/resource-manager/Microsoft.DocumentDB/preview/${version}/${githubResourceName}.json`;
 const outputDir = path.join(__dirname, `../../src/Utils/arm/generatedClients/${deResourceName}`);
 
@@ -117,9 +117,9 @@ const propertyToType = (property: Property, prop: string, required: boolean) =>
     if (property.allOf) {
       outputTypes.push(`
       /* ${property.description || "undocumented"} */
-      ${property.readOnly ? "readonly " : ""}${prop}${
+      ${property.readOnly ? "readonly " : ""}${prop}${required ? "" : "?"}: ${property.allOf
-        required ? "" : "?"
+        .map((allof: { $ref: string }) => refToType(allof.$ref))
-      }: ${property.allOf.map((allof: { $ref: string }) => refToType(allof.$ref)).join(" & ")}`);
+        .join(" & ")}`);
     } else if (property.$ref) {
       const type = refToType(property.$ref);
       outputTypes.push(`
@@ -142,8 +142,8 @@ const propertyToType = (property: Property, prop: string, required: boolean) =>
       outputTypes.push(`
           /* ${property.description || "undocumented"} */
           ${property.readOnly ? "readonly " : ""}${prop}${required ? "" : "?"}: ${property.enum
-        .map((v: string) => `"${v}"`)
+            .map((v: string) => `"${v}"`)
-        .join(" | ")}
+            .join(" | ")}
           `);
     } else {
       if (property.type === undefined) {
@@ -153,8 +153,8 @@ const propertyToType = (property: Property, prop: string, required: boolean) =>
       outputTypes.push(`
           /* ${property.description || "undocumented"} */
           ${property.readOnly ? "readonly " : ""}${prop}${required ? "" : "?"}: ${
-        propertyMap[property.type] ? propertyMap[property.type] : property.type
+            propertyMap[property.type] ? propertyMap[property.type] : property.type
-      }`);
+          }`);
     }
   }
 };
@@ -247,7 +247,7 @@ async function main() {
         const operation = schema.paths[path][method];
         const [, methodName] = operation.operationId.split("_");
         const bodyParameter = operation.parameters.find(
-          (parameter: { in: string; required: boolean }) => parameter.in === "body" && parameter.required === true
+          (parameter: { in: string; required: boolean }) => parameter.in === "body" && parameter.required === true,
         );
         outputClient.push(`
           /* ${operation.description || "undocumented"} */
@@ -259,8 +259,8 @@ async function main() {
           ) : Promise<${responseType(operation, "Types")}> {
             const path = \`${path.replace(/{/g, "${")}\`
             return armRequest({ host: configContext.ARM_ENDPOINT, path, method: "${method.toLocaleUpperCase()}", apiVersion, ${
-          bodyParameter ? "body" : ""
+              bodyParameter ? "body" : ""
-        } })
+            } })
           }
           `);
       }