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;

src/Contracts/ViewModels.ts

@@ -106,8 +106,8 @@ export interface CollectionBase extends TreeNode {
   self: string;
   rawDataModel: DataModels.Collection;
   partitionKey: DataModels.PartitionKey;
-  partitionKeyProperty: string;
+  partitionKeyProperties: string[];
-  partitionKeyPropertyHeader: string;
+  partitionKeyPropertyHeaders: string[];
   id: ko.Observable<string>;
   selectedSubnodeKind: ko.Observable<CollectionTabKind>;
   children: ko.ObservableArray<TreeNode>;

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

@@ -65,8 +65,8 @@ export class SubSettingsComponent extends React.Component<SubSettingsComponentPr
   constructor(props: SubSettingsComponentProps) {
     super(props);
     this.geospatialVisible = userContext.apiType === "SQL";
-    this.partitionKeyValue = "/" + this.props.collection.partitionKeyProperty;
     this.partitionKeyName = userContext.apiType === "Mongo" ? "Shard key" : "Partition key";
+    this.partitionKeyValue = this.getPartitionKeyValue();
   }
 
   componentDidMount(): void {
@@ -291,6 +291,14 @@ export class SubSettingsComponent extends React.Component<SubSettingsComponentPr
     );
   };
 
+  private getPartitionKeyValue = (): string => {
+    if (userContext.apiType === "Mongo") {
+      return this.props.collection.partitionKeyProperties?.[0] || "";
+    }
+
+    return (this.props.collection.partitionKeyProperties || []).map((property) => "/" + property).join(", ");
+  };
+
   private getPartitionKeyComponent = (): JSX.Element => (
     <Stack {...titleAndInputStackProps}>
       {this.getPartitionKeyVisible() && (
@@ -310,7 +318,8 @@ export class SubSettingsComponent extends React.Component<SubSettingsComponentPr
     if (
       userContext.apiType === "Cassandra" ||
       userContext.apiType === "Tables" ||
-      !this.props.collection.partitionKeyProperty ||
+      !this.props.collection.partitionKeyProperties ||
+      this.props.collection.partitionKeyProperties.length === 0 ||
       (userContext.apiType === "Mongo" && this.props.collection.partitionKey.systemKey)
     ) {
       return false;

src/Explorer/Controls/Settings/TestUtils.tsx

@@ -39,7 +39,7 @@ export const collection = ({
     kind: "hash",
     version: 2,
   },
-  partitionKeyProperty: "partitionKey",
+  partitionKeyProperties: ["partitionKey"],
   readSettings: () => {
     return;
   },

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

@@ -1669,7 +1669,9 @@ exports[`SettingsComponent renders 1`] = `
                 "paths": Array [],
                 "version": 2,
               },
-              "partitionKeyProperty": "partitionKey",
+              "partitionKeyProperties": Array [
+                "partitionKey",
+              ],
               "readSettings": [Function],
               "uniqueKeyPolicy": Object {},
               "usageSizeInKB": [Function],
@@ -3348,7 +3350,9 @@ exports[`SettingsComponent renders 1`] = `
                 "paths": Array [],
                 "version": 2,
               },
-              "partitionKeyProperty": "partitionKey",
+              "partitionKeyProperties": Array [
+                "partitionKey",
+              ],
               "readSettings": [Function],
               "uniqueKeyPolicy": Object {},
               "usageSizeInKB": [Function],

src/Explorer/OpenActions/OpenActions.tsx

@@ -8,23 +8,23 @@ import { CassandraAddCollectionPane } from "../Panes/CassandraAddCollectionPane/
 import { SettingsPane } from "../Panes/SettingsPane/SettingsPane";
 import { CassandraAPIDataClient } from "../Tables/TableDataClient";
 
-function generateQueryText(action: ActionContracts.OpenQueryTab, partitionKeyProperty: string): string {
+function generateQueryText(action: ActionContracts.OpenQueryTab, partitionKeyProperties: string[]): string {
   if (!action.query) {
     return "SELECT * FROM c";
   } else if (action.query.text) {
     return action.query.text;
-  } else if (!!action.query.partitionKeys && action.query.partitionKeys.length > 0) {
+  } else if (action.query.partitionKeys?.length > 0 && partitionKeyProperties?.length > 0) {
     let query = "SELECT * FROM c WHERE";
     for (let i = 0; i < action.query.partitionKeys.length; i++) {
       const partitionKey = action.query.partitionKeys[i];
       if (!partitionKey) {
         // null partition key case
-        query = query.concat(` c.${partitionKeyProperty} = ${action.query.partitionKeys[i]}`);
+        query = query.concat(` c.${partitionKeyProperties[i]} = ${action.query.partitionKeys[i]}`);
       } else if (typeof partitionKey !== "string") {
         // Undefined partition key case
-        query = query.concat(` NOT IS_DEFINED(c.${partitionKeyProperty})`);
+        query = query.concat(` NOT IS_DEFINED(c.${partitionKeyProperties[i]})`);
       } else {
-        query = query.concat(` c.${partitionKeyProperty} = "${action.query.partitionKeys[i]}"`);
+        query = query.concat(` c.${partitionKeyProperties[i]} = "${action.query.partitionKeys[i]}"`);
       }
       if (i !== action.query.partitionKeys.length - 1) {
         query = query.concat(" OR");
@@ -109,7 +109,7 @@ function openCollectionTab(
           collection.onNewQueryClick(
             collection,
             undefined,
-            generateQueryText(action as ActionContracts.OpenQueryTab, collection.partitionKeyProperty)
+            generateQueryText(action as ActionContracts.OpenQueryTab, collection.partitionKeyProperties)
           );
           break;
         }

src/Explorer/Tables/TableEntityProcessor.ts

@@ -126,7 +126,7 @@ export function convertEntitiesToDocuments(
       };
       if (collection.partitionKey) {
         document["partitionKey"] = collection.partitionKey;
-        document[collection.partitionKeyProperty] = entity.PartitionKey._;
+        document[collection.partitionKeyProperties[0]] = entity.PartitionKey._;
         document["partitionKeyValue"] = entity.PartitionKey._;
       }
       for (var property in entity) {

src/Explorer/Tabs/ConflictsTab.ts

@@ -74,7 +74,7 @@ export default class ConflictsTab extends TabsBase {
     this.partitionKey = options.partitionKey || (this.collection && this.collection.partitionKey);
     this.conflictIds = options.conflictIds;
     this.partitionKeyPropertyHeader =
-      (this.collection && this.collection.partitionKeyPropertyHeader) || this._getPartitionKeyPropertyHeader();
+      this.collection?.partitionKeyPropertyHeaders?.[0] || this._getPartitionKeyPropertyHeader();
     this.partitionKeyProperty = !!this.partitionKeyPropertyHeader
       ? this.partitionKeyPropertyHeader.replace(/[/]+/g, ".").substr(1).replace(/[']+/g, "")
       : null;

src/Explorer/Tabs/DocumentsTab.html

@@ -143,16 +143,18 @@
               <tr>
                 <th class="documentsGridHeader" data-bind="text: idHeader" tabindex="0"></th>
                 <!-- ko if: showPartitionKey -->
+                <!-- ko foreach: partitionKeyPropertyHeaders -->
                 <th
                   class="documentsGridHeader documentsGridPartition evenlySpacedHeader"
                   data-bind="
                                               attr: {
-                                                  title: partitionKeyPropertyHeader
+                                                  title: $data
                                               },
-                                              text: partitionKeyPropertyHeader"
+                                              text: $data"
                   tabindex="0"
                 ></th>
                 <!-- /ko -->
+                <!-- /ko -->
                 <th
                   class="refreshColHeader"
                   role="button"
@@ -182,13 +184,13 @@
                 tabindex="0"
               >
                 <td class="tabdocumentsGridElement"><a data-bind="text: $data.id, attr: { title: $data.id }"></a></td>
-                <!-- ko if: $data.partitionKeyProperty -->
+                <!-- ko if: $data.partitionKeyProperties -->
-                <td class="tabdocumentsGridElement" colspan="2">
+                <!-- ko foreach: $data.stringPartitionKeyValues -->
-                  <a
+                <td class="tabdocumentsGridElement" data-bind="colspan: $parent.stringPartitionKeyValues.length + 1">
-                    data-bind="text: $data.stringPartitionKeyValue, attr: { title: $data.stringPartitionKeyValue }"
+                  <a data-bind="text: $data, attr: { title: $data }"></a>
-                  ></a>
                 </td>
                 <!-- /ko -->
+                <!-- /ko -->
               </tr>
               <!-- /ko -->
             </tbody>

src/Explorer/Tabs/DocumentsTab.ts

@@ -50,7 +50,7 @@ export default class DocumentsTab extends TabsBase {
   public editorState: ko.Observable<ViewModels.DocumentExplorerState>;
   public newDocumentButton: ViewModels.Button;
   public saveNewDocumentButton: ViewModels.Button;
-  public saveExisitingDocumentButton: ViewModels.Button;
+  public saveExistingDocumentButton: ViewModels.Button;
   public discardNewDocumentChangesButton: ViewModels.Button;
   public discardExisitingDocumentChangesButton: ViewModels.Button;
   public deleteExisitingDocumentButton: ViewModels.Button;
@@ -65,8 +65,8 @@ export default class DocumentsTab extends TabsBase {
 
   // TODO need to refactor
   public partitionKey: DataModels.PartitionKey;
-  public partitionKeyPropertyHeader: string;
+  public partitionKeyPropertyHeaders: string[];
-  public partitionKeyProperty: string;
+  public partitionKeyProperties: string[];
   public documentIds: ko.ObservableArray<DocumentId>;
 
   private _documentsIterator: QueryIterator<ItemDefinition & Resource>;
@@ -90,11 +90,10 @@ export default class DocumentsTab extends TabsBase {
     this._resourceTokenPartitionKey = options.resourceTokenPartitionKey;
     this.documentIds = options.documentIds;
 
-    this.partitionKeyPropertyHeader =
+    this.partitionKeyPropertyHeaders = this.collection?.partitionKeyPropertyHeaders || this.partitionKey?.paths;
-      (this.collection && this.collection.partitionKeyPropertyHeader) || this._getPartitionKeyPropertyHeader();
+    this.partitionKeyProperties = this.partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader) =>
-    this.partitionKeyProperty = !!this.partitionKeyPropertyHeader
+      partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, "")
-      ? this.partitionKeyPropertyHeader.replace(/[/]+/g, ".").substr(1).replace(/[']+/g, "")
+    );
-      : null;
 
     this.isFilterExpanded = ko.observable<boolean>(false);
     this.isFilterCreated = ko.observable<boolean>(true);
@@ -227,7 +226,7 @@ export default class DocumentsTab extends TabsBase {
       }),
     };
 
-    this.saveExisitingDocumentButton = {
+    this.saveExistingDocumentButton = {
       enabled: ko.computed<boolean>(() => {
         switch (this.editorState()) {
           case ViewModels.DocumentExplorerState.exisitingDocumentDirtyValid:
@@ -445,8 +444,7 @@ export default class DocumentsTab extends TabsBase {
             savedDocument,
             this.partitionKey as PartitionKeyDefinition
           );
-          const partitionKeyValue = partitionKeyValueArray && partitionKeyValueArray[0];
+          let id = new DocumentId(this, savedDocument, partitionKeyValueArray);
-          let id = new DocumentId(this, savedDocument, partitionKeyValue);
           let ids = this.documentIds();
           ids.push(id);
 
@@ -489,14 +487,12 @@ export default class DocumentsTab extends TabsBase {
     return Q();
   };
 
-  public onSaveExisitingDocumentClick = (): Promise<any> => {
+  public onSaveExistingDocumentClick = (): Promise<any> => {
     const selectedDocumentId = this.selectedDocumentId();
     const documentContent = JSON.parse(this.selectedDocumentContent());
 
     const partitionKeyValueArray = extractPartitionKey(documentContent, this.partitionKey as PartitionKeyDefinition);
-    const partitionKeyValue = partitionKeyValueArray && partitionKeyValueArray[0];
+    selectedDocumentId.partitionKeyValue = partitionKeyValueArray;
-
-    selectedDocumentId.partitionKeyValue = partitionKeyValue;
 
     this.isExecutionError(false);
     const startKey: number = TelemetryProcessor.traceStart(Action.UpdateDocument, {
@@ -800,7 +796,7 @@ export default class DocumentsTab extends TabsBase {
   }
 
   public buildQuery(filter: string): string {
-    return QueryUtils.buildDocumentsQuery(filter, this.partitionKeyProperty, this.partitionKey);
+    return QueryUtils.buildDocumentsQuery(filter, this.partitionKeyProperties, this.partitionKey);
   }
 
   protected getTabsButtons(): CommandButtonComponentProps[] {
@@ -844,16 +840,16 @@ export default class DocumentsTab extends TabsBase {
       });
     }
 
-    if (this.saveExisitingDocumentButton.visible()) {
+    if (this.saveExistingDocumentButton.visible()) {
       const label = "Update";
       buttons.push({
         iconSrc: SaveIcon,
         iconAlt: label,
-        onCommandClick: this.onSaveExisitingDocumentClick,
+        onCommandClick: this.onSaveExistingDocumentClick,
         commandButtonLabel: label,
         ariaLabel: label,
         hasPopup: false,
-        disabled: !this.saveExisitingDocumentButton.enabled(),
+        disabled: !this.saveExistingDocumentButton.enabled(),
       });
     }
 
@@ -899,8 +895,8 @@ export default class DocumentsTab extends TabsBase {
         this.saveNewDocumentButton.enabled,
         this.discardNewDocumentChangesButton.visible,
         this.discardNewDocumentChangesButton.enabled,
-        this.saveExisitingDocumentButton.visible,
+        this.saveExistingDocumentButton.visible,
-        this.saveExisitingDocumentButton.enabled,
+        this.saveExistingDocumentButton.enabled,
         this.discardExisitingDocumentChangesButton.visible,
         this.discardExisitingDocumentChangesButton.enabled,
         this.deleteExisitingDocumentButton.visible,
@@ -910,16 +906,6 @@ export default class DocumentsTab extends TabsBase {
     this.updateNavbarWithTabsButtons();
   }
 
-  private _getPartitionKeyPropertyHeader(): string {
-    return (
-      (this.partitionKey &&
-        this.partitionKey.paths &&
-        this.partitionKey.paths.length > 0 &&
-        this.partitionKey.paths[0]) ||
-      null
-    );
-  }
-
   public static _createUploadButton(container: Explorer): CommandButtonComponentProps {
     const label = "Upload Item";
     return {

src/Explorer/Tabs/MongoDocumentsTab.ts

@@ -29,16 +29,20 @@ export default class MongoDocumentsTab extends DocumentsTab {
     super(options);
     this.lastFilterContents = ko.observableArray<string>(['{"id":"foo"}', "{ qty: { $gte: 20 } }"]);
 
-    if (this.partitionKeyProperty && ~this.partitionKeyProperty.indexOf(`"`)) {
+    this.partitionKeyProperties = this.partitionKeyProperties?.map((partitionKeyProperty, i) => {
-      this.partitionKeyProperty = this.partitionKeyProperty.replace(/["]+/g, "");
+      if (partitionKeyProperty && ~partitionKeyProperty.indexOf(`"`)) {
+        partitionKeyProperty = partitionKeyProperty.replace(/["]+/g, "");
       }
 
-    if (this.partitionKeyProperty && this.partitionKeyProperty.indexOf("$v") > -1) {
+      if (partitionKeyProperty && partitionKeyProperty.indexOf("$v") > -1) {
         // From $v.shard.$v.key.$v > shard.key
-      this.partitionKeyProperty = this.partitionKeyProperty.replace(/.\$v/g, "").replace(/\$v./g, "");
+        partitionKeyProperty = partitionKeyProperty.replace(/.\$v/g, "").replace(/\$v./g, "");
-      this.partitionKeyPropertyHeader = "/" + this.partitionKeyProperty;
+        this.partitionKeyPropertyHeaders[i] = "/" + partitionKeyProperty;
       }
 
+      return partitionKeyProperty;
+    });
+
     this.isFilterExpanded = ko.observable<boolean>(true);
     super.buildCommandBarOptions.bind(this);
     super.buildCommandBarOptions();
@@ -52,12 +56,9 @@ export default class MongoDocumentsTab extends DocumentsTab {
       tabTitle: this.tabTitle(),
     });
 
-    if (
+    const partitionKeyProperty = this.partitionKeyProperties?.[0];
-      this.partitionKeyProperty &&
+    if (partitionKeyProperty !== "_id" && !this._hasShardKeySpecified(documentContent)) {
-      this.partitionKeyProperty !== "_id" &&
+      const message = `The document is lacking the shard property: ${partitionKeyProperty}`;
-      !this._hasShardKeySpecified(documentContent)
-    ) {
-      const message = `The document is lacking the shard property: ${this.partitionKeyProperty}`;
       this.displayedError(message);
       let that = this;
       setTimeout(() => {
@@ -79,7 +80,12 @@ export default class MongoDocumentsTab extends DocumentsTab {
 
     this.isExecutionError(false);
     this.isExecuting(true);
-    return createDocument(this.collection.databaseId, this.collection, this.partitionKeyProperty, documentContent)
+    return createDocument(
+      this.collection.databaseId,
+      this.collection,
+      this.partitionKeyProperties?.[0],
+      documentContent
+    )
       .then(
         (savedDocument: any) => {
           let partitionKeyArray = extractPartitionKey(
@@ -87,9 +93,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
             this._getPartitionKeyDefinition() as PartitionKeyDefinition
           );
 
-          let partitionKeyValue = partitionKeyArray && partitionKeyArray[0];
+          let id = new ObjectId(this, savedDocument, partitionKeyArray);
-
-          let id = new ObjectId(this, savedDocument, partitionKeyValue);
           let ids = this.documentIds();
           ids.push(id);
           delete savedDocument._self;
@@ -128,7 +132,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
       .finally(() => this.isExecuting(false));
   };
 
-  public onSaveExisitingDocumentClick = (): Promise<any> => {
+  public onSaveExistingDocumentClick = (): Promise<any> => {
     const selectedDocumentId = this.selectedDocumentId();
     const documentContent = this.selectedDocumentContent();
     this.isExecutionError(false);
@@ -151,9 +155,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
                 this._getPartitionKeyDefinition() as PartitionKeyDefinition
               );
 
-              let partitionKeyValue = partitionKeyArray && partitionKeyArray[0];
+              const id = new ObjectId(this, updatedDocument, partitionKeyArray);
-
-              const id = new ObjectId(this, updatedDocument, partitionKeyValue);
               documentId.id(id.id());
             }
           });
@@ -214,7 +216,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
             })
             .map((rawDocument: any) => {
               const partitionKeyValue = rawDocument._partitionKeyValue;
-              return new DocumentId(this, rawDocument, partitionKeyValue);
+              return new DocumentId(this, rawDocument, [partitionKeyValue]);
             });
 
           const merged = currentDocuments.concat(nextDocumentIds);
@@ -303,7 +305,7 @@ export default class MongoDocumentsTab extends DocumentsTab {
       // Convert BsonSchema2 to /path format
       partitionKey = {
         kind: partitionKey.kind,
-        paths: ["/" + this.partitionKeyProperty.replace(/\./g, "/")],
+        paths: ["/" + this.partitionKeyProperties?.[0].replace(/\./g, "/")],
         version: partitionKey.version,
       };
     }

src/Explorer/Tree/Collection.test.ts

@@ -37,7 +37,8 @@ describe("Collection", () => {
         version: 2,
       });
       collection = generateMockCollectionWithDataModel(collectionsDataModel);
-      expect(collection.partitionKeyProperty).toBe("somePartitionKey.anotherPartitionKey");
+      expect(collection.partitionKeyProperties.length).toBe(1);
+      expect(collection.partitionKeyProperties[0]).toBe("somePartitionKey.anotherPartitionKey");
     });
 
     it("should strip out forward slashes from single partition key paths", () => {
@@ -47,7 +48,8 @@ describe("Collection", () => {
         version: 2,
       });
       collection = generateMockCollectionWithDataModel(collectionsDataModel);
-      expect(collection.partitionKeyProperty).toBe("somePartitionKey");
+      expect(collection.partitionKeyProperties.length).toBe(1);
+      expect(collection.partitionKeyProperties[0]).toBe("somePartitionKey");
     });
   });
 
@@ -61,7 +63,8 @@ describe("Collection", () => {
         version: 2,
       });
       collection = generateMockCollectionWithDataModel(collectionsDataModel);
-      expect(collection.partitionKeyPropertyHeader).toBe("/somePartitionKey/anotherPartitionKey");
+      expect(collection.partitionKeyPropertyHeaders.length).toBe(1);
+      expect(collection.partitionKeyPropertyHeaders[0]).toBe("/somePartitionKey/anotherPartitionKey");
     });
 
     it("should preserve forward slash on a single partition key", () => {
@@ -71,7 +74,8 @@ describe("Collection", () => {
         version: 2,
       });
       collection = generateMockCollectionWithDataModel(collectionsDataModel);
-      expect(collection.partitionKeyPropertyHeader).toBe("/somePartitionKey");
+      expect(collection.partitionKeyPropertyHeaders.length).toBe(1);
+      expect(collection.partitionKeyPropertyHeaders[0]).toBe("/somePartitionKey");
     });
 
     it("should be null if there is no partition key", () => {
@@ -81,7 +85,7 @@ describe("Collection", () => {
         kind: "Hash",
       });
       collection = generateMockCollectionWithDataModel(collectionsDataModel);
-      expect(collection.partitionKeyPropertyHeader).toBeNull();
+      expect(collection.partitionKeyPropertyHeaders.length).toBe(0);
     });
   });
 });

src/Explorer/Tree/Collection.ts

@@ -50,8 +50,8 @@ export default class Collection implements ViewModels.Collection {
   public rid: string;
   public databaseId: string;
   public partitionKey: DataModels.PartitionKey;
-  public partitionKeyPropertyHeader: string;
+  public partitionKeyPropertyHeaders: string[];
-  public partitionKeyProperty: string;
+  public partitionKeyProperties: string[];
   public id: ko.Observable<string>;
   public defaultTtl: ko.Observable<number>;
   public indexingPolicy: ko.Observable<DataModels.IndexingPolicy>;
@@ -120,31 +120,25 @@ export default class Collection implements ViewModels.Collection {
     this.requestSchema = data.requestSchema;
     this.geospatialConfig = ko.observable(data.geospatialConfig);
 
+    this.partitionKeyPropertyHeaders = this.partitionKey?.paths;
+    this.partitionKeyProperties = this.partitionKeyPropertyHeaders?.map((partitionKeyPropertyHeader, i) => {
       // TODO fix this to only replace non-excaped single quotes
-    this.partitionKeyProperty =
+      let partitionKeyProperty = partitionKeyPropertyHeader.replace(/[/]+/g, ".").substring(1).replace(/[']+/g, "");
-      (this.partitionKey &&
-        this.partitionKey.paths &&
-        this.partitionKey.paths.length &&
-        this.partitionKey.paths.length > 0 &&
-        this.partitionKey.paths[0].replace(/[/]+/g, ".").substr(1).replace(/[']+/g, "")) ||
-      null;
-    this.partitionKeyPropertyHeader =
-      (this.partitionKey &&
-        this.partitionKey.paths &&
-        this.partitionKey.paths.length > 0 &&
-        this.partitionKey.paths[0]) ||
-      null;
 
-    if (userContext.apiType === "Mongo" && this.partitionKeyProperty && ~this.partitionKeyProperty.indexOf(`"`)) {
+      if (userContext.apiType === "Mongo" && partitionKeyProperty) {
-      this.partitionKeyProperty = this.partitionKeyProperty.replace(/["]+/g, "");
+        if (~partitionKeyProperty.indexOf(`"`)) {
+          partitionKeyProperty = partitionKeyProperty.replace(/["]+/g, "");
         }
-
         // TODO #10738269 : Add this logic in a derived class for Mongo
-    if (userContext.apiType === "Mongo" && this.partitionKeyProperty && this.partitionKeyProperty.indexOf("$v") > -1) {
+        if (partitionKeyProperty.indexOf("$v") > -1) {
           // From $v.shard.$v.key.$v > shard.key
-      this.partitionKeyProperty = this.partitionKeyProperty.replace(/.\$v/g, "").replace(/\$v./g, "");
+          partitionKeyProperty = partitionKeyProperty.replace(/.\$v/g, "").replace(/\$v./g, "");
-      this.partitionKeyPropertyHeader = "/" + this.partitionKeyProperty;
+          this.partitionKeyPropertyHeaders[i] = partitionKeyProperty;
         }
+      }
+
+      return partitionKeyProperty;
+    });
 
     this.documentIds = ko.observableArray<DocumentId>([]);
     this.isCollectionExpanded = ko.observable<boolean>(false);
@@ -471,7 +465,7 @@ export default class Collection implements ViewModels.Collection {
 
         collection: this,
         masterKey: userContext.masterKey || "",
-        collectionPartitionKeyProperty: this.partitionKeyProperty,
+        collectionPartitionKeyProperty: this.partitionKeyProperties?.[0],
         collectionId: this.id(),
         databaseId: this.databaseId,
         isTabsContentExpanded: this.container.isTabsContentExpanded,
@@ -710,7 +704,7 @@ export default class Collection implements ViewModels.Collection {
       tabPath: "",
       collection: this,
       masterKey: userContext.masterKey || "",
-      collectionPartitionKeyProperty: this.partitionKeyProperty,
+      collectionPartitionKeyProperty: this.partitionKeyProperties?.[0],
       collectionId: this.id(),
       databaseId: this.databaseId,
       isTabsContentExpanded: this.container.isTabsContentExpanded,

src/Explorer/Tree/ConflictId.ts

@@ -150,7 +150,7 @@ export default class ConflictId {
       partitionKeyValueResolved
     );
 
-    documentId.partitionKeyProperty = this.partitionKeyProperty;
+    documentId.partitionKeyProperties = [this.partitionKeyProperty];
     documentId.partitionKey = this.partitionKey;
 
     return documentId;

src/Explorer/Tree/DocumentId.ts

@@ -9,21 +9,21 @@ export default class DocumentId {
   public self: string;
   public ts: string;
   public id: ko.Observable<string>;
-  public partitionKeyProperty: string;
+  public partitionKeyProperties: string[];
   public partitionKey: DataModels.PartitionKey;
-  public partitionKeyValue: any;
+  public partitionKeyValue: any[];
-  public stringPartitionKeyValue: string;
+  public stringPartitionKeyValues: string[];
   public isDirty: ko.Observable<boolean>;
 
-  constructor(container: DocumentsTab, data: any, partitionKeyValue: any) {
+  constructor(container: DocumentsTab, data: any, partitionKeyValue: any[]) {
     this.container = container;
     this.self = data._self;
     this.rid = data._rid;
     this.ts = data._ts;
     this.partitionKeyValue = partitionKeyValue;
-    this.partitionKeyProperty = container && container.partitionKeyProperty;
+    this.partitionKeyProperties = container?.partitionKeyProperties;
     this.partitionKey = container && container.partitionKey;
-    this.stringPartitionKeyValue = this.getPartitionKeyValueAsString();
+    this.stringPartitionKeyValues = this.getPartitionKeyValueAsString();
     this.id = ko.observable(data.id);
     this.isDirty = ko.observable(false);
   }
@@ -46,19 +46,19 @@ export default class DocumentId {
   }
 
   public partitionKeyHeader(): Object {
-    if (!this.partitionKeyProperty) {
+    if (!this.partitionKeyProperties || this.partitionKeyProperties.length === 0) {
       return undefined;
     }
 
-    if (this.partitionKeyValue === undefined) {
+    if (!this.partitionKeyValue || this.partitionKeyValue.length === 0) {
       return [{}];
     }
 
     return [this.partitionKeyValue];
   }
 
-  public getPartitionKeyValueAsString(): string {
+  public getPartitionKeyValueAsString(): string[] {
-    const partitionKeyValue: any = this.partitionKeyValue;
+    return this.partitionKeyValue?.map((partitionKeyValue) => {
       const typeOfPartitionKeyValue: string = typeof partitionKeyValue;
 
       if (
@@ -74,6 +74,7 @@ export default class DocumentId {
       }
 
       return JSON.stringify(partitionKeyValue);
+    });
   }
 
   public async loadDocument(): Promise<void> {

src/Explorer/Tree/ResourceTokenCollection.ts

@@ -22,8 +22,8 @@ export default class ResourceTokenCollection implements ViewModels.CollectionBas
   public rid: string;
   public rawDataModel: DataModels.Collection;
   public partitionKey: DataModels.PartitionKey;
-  public partitionKeyProperty: string;
+  public partitionKeyProperties: string[];
-  public partitionKeyPropertyHeader: string;
+  public partitionKeyPropertyHeaders: string[];
   public id: ko.Observable<string>;
   public children: ko.ObservableArray<ViewModels.TreeNode>;
   public selectedSubnodeKind: ko.Observable<ViewModels.CollectionTabKind>;

src/Utils/QueryUtils.ts

@@ -3,14 +3,15 @@ import * as ViewModels from "../Contracts/ViewModels";
 
 export function buildDocumentsQuery(
   filter: string,
-  partitionKeyProperty: string,
+  partitionKeyProperties: string[],
   partitionKey: DataModels.PartitionKey
 ): string {
-  let query = partitionKeyProperty
+  let query =
-    ? `select c.id, c._self, c._rid, c._ts, ${buildDocumentsQueryPartitionProjections(
+    partitionKeyProperties && partitionKeyProperties.length > 0
+      ? `select c.id, c._self, c._rid, c._ts, [${buildDocumentsQueryPartitionProjections(
           "c",
           partitionKey
-      )} as _partitionKeyValue from c`
+        )}] as _partitionKeyValue from c`
       : `select c.id, c._self, c._rid, c._ts from c`;
 
   if (filter) {