Add support for multi-partition key (#1252)

src/Common/MongoProxyClient.test.ts

@@ -25,12 +25,12 @@ const fetchMock = () => {
   });
 };
 
-const partitionKeyProperty = "pk";
+const partitionKeyProperties = ["pk"];
 
 const collection = {
   id: () => "testCollection",
   rid: "testCollectionrid",
-  partitionKeyProperty,
+  partitionKeyProperties,
   partitionKey: {
     paths: ["/pk"],
     kind: "Hash",
@@ -41,7 +41,7 @@ const collection = {
 const documentId = ({
   partitionKeyHeader: () => "[]",
   self: "db/testDB/db/testCollection/docs/testId",
-  partitionKeyProperty,
+  partitionKeyProperties,
   partitionKey: {
     paths: ["/pk"],
     kind: "Hash",

src/Common/MongoProxyClient.ts

@@ -76,7 +76,7 @@ export function queryDocuments(
     dba: databaseAccount.name,
     pk:
       collection && collection.partitionKey && !collection.partitionKey.systemKey
-        ? collection.partitionKeyProperty
+        ? collection.partitionKeyProperties?.[0]
         : "",
   };
 
@@ -139,7 +139,7 @@ export function readDocument(
     dba: databaseAccount.name,
     pk:
       documentId && documentId.partitionKey && !documentId.partitionKey.systemKey
-        ? documentId.partitionKeyProperty
+        ? documentId.partitionKeyProperties?.[0]
         : "",
   };
 
@@ -225,7 +225,7 @@ export function updateDocument(
     dba: databaseAccount.name,
     pk:
       documentId && documentId.partitionKey && !documentId.partitionKey.systemKey
-        ? documentId.partitionKeyProperty
+        ? documentId.partitionKeyProperties?.[0]
         : "",
   };
   const endpoint = getFeatureEndpointOrDefault("updateDocument");
@@ -266,7 +266,7 @@ export function deleteDocument(databaseId: string, collection: Collection, docum
     dba: databaseAccount.name,
     pk:
       documentId && documentId.partitionKey && !documentId.partitionKey.systemKey
-        ? documentId.partitionKeyProperty
+        ? documentId.partitionKeyProperties?.[0]
         : "",
   };
   const endpoint = getFeatureEndpointOrDefault("deleteDocument");

src/Common/QueriesClient.ts

@@ -149,10 +149,10 @@ export class QueriesClient {
     const documentId = new DocumentId(
       {
         partitionKey: QueriesClient.PartitionKey,
-        partitionKeyProperty: "id",
+        partitionKeyProperties: ["id"],
       } as DocumentsTab,
       query,
-      query.queryName
+      [query.queryName]
     ); // TODO: Remove DocumentId's dependency on DocumentsTab
     const options: any = { partitionKey: query.resourceId };
     return deleteDocument(queriesCollection, documentId)

src/Common/dataAccess/readDocument.ts

@@ -1,21 +1,28 @@
-import { Item } from "@azure/cosmos";
+import { Item, RequestOptions } from "@azure/cosmos";
 import { CollectionBase } from "../../Contracts/ViewModels";
+import DocumentId from "../../Explorer/Tree/DocumentId";
+import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
+import { HttpHeaders } from "../Constants";
 import { client } from "../CosmosClient";
 import { getEntityName } from "../DocumentUtility";
 import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import DocumentId from "../../Explorer/Tree/DocumentId";
 
 export const readDocument = async (collection: CollectionBase, documentId: DocumentId): Promise<Item> => {
   const entityName = getEntityName();
   const clearMessage = logConsoleProgress(`Reading ${entityName} ${documentId.id()}`);
 
   try {
+    const options: RequestOptions =
+      documentId.partitionKey.kind === "MultiHash"
+        ? {
+            [HttpHeaders.partitionKey]: documentId.partitionKeyValue,
+          }
+        : {};
     const response = await client()
       .database(collection.databaseId)
       .container(collection.id())
       .item(documentId.id(), documentId.partitionKeyValue)
-      .read();
+      .read(options);
 
     return response?.resource;
   } catch (error) {

src/Common/dataAccess/updateDocument.ts

@@ -1,10 +1,11 @@
+import { Item, RequestOptions } from "@azure/cosmos";
+import { HttpHeaders } from "Common/Constants";
 import { CollectionBase } from "../../Contracts/ViewModels";
-import { Item } from "@azure/cosmos";
+import DocumentId from "../../Explorer/Tree/DocumentId";
+import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
 import { client } from "../CosmosClient";
 import { getEntityName } from "../DocumentUtility";
 import { handleError } from "../ErrorHandlingUtils";
-import { logConsoleInfo, logConsoleProgress } from "../../Utils/NotificationConsoleUtils";
-import DocumentId from "../../Explorer/Tree/DocumentId";
 
 export const updateDocument = async (
   collection: CollectionBase,
@@ -15,11 +16,17 @@ export const updateDocument = async (
   const clearMessage = logConsoleProgress(`Updating ${entityName} ${documentId.id()}`);
 
   try {
+    const options: RequestOptions =
+      documentId.partitionKey.kind === "MultiHash"
+        ? {
+            [HttpHeaders.partitionKey]: documentId.partitionKeyValue,
+          }
+        : {};
     const response = await client()
       .database(collection.databaseId)
       .container(collection.id())
       .item(documentId.id(), documentId.partitionKeyValue)
-      .replace(newDocument);
+      .replace(newDocument, options);
 
     logConsoleInfo(`Successfully updated ${entityName} ${documentId.id()}`);
     return response?.resource;