Fix for Mongo shard key loaded via ARM (#657)

src/Contracts/DataModels.ts

@@ -121,6 +121,10 @@ export interface ISchemaRequest {
 }
 
 export interface Collection extends Resource {
+  // Only in Mongo collections loaded via ARM
+  shardKey?: {
+    [key: string]: string;
+  };
   defaultTtl?: number;
   indexingPolicy?: IndexingPolicy;
   partitionKey?: PartitionKey;

src/Explorer/Tabs/MongoDocumentsTab.ts

@@ -1,14 +1,9 @@
-import * as Constants from "../../Common/Constants";
-import * as DataModels from "../../Contracts/DataModels";
+import { extractPartitionKey, PartitionKeyDefinition } from "@azure/cosmos";
 import * as ko from "knockout";
-import * as ViewModels from "../../Contracts/ViewModels";
-import DocumentId from "../Tree/DocumentId";
-import DocumentsTab from "./DocumentsTab";
-import MongoUtility from "../../Common/MongoUtility";
-import ObjectId from "../Tree/ObjectId";
 import Q from "q";
-import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
-import { Action } from "../../Shared/Telemetry/TelemetryConstants";
+import * as Constants from "../../Common/Constants";
+import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
+import * as Logger from "../../Common/Logger";
 import {
   createDocument,
   deleteDocument,
@@ -16,10 +11,14 @@ import {
   readDocument,
   updateDocument,
 } from "../../Common/MongoProxyClient";
-import { extractPartitionKey } from "@azure/cosmos";
-import * as Logger from "../../Common/Logger";
-import { PartitionKeyDefinition } from "@azure/cosmos";
-import { getErrorMessage, getErrorStack } from "../../Common/ErrorHandlingUtils";
+import MongoUtility from "../../Common/MongoUtility";
+import * as DataModels from "../../Contracts/DataModels";
+import * as ViewModels from "../../Contracts/ViewModels";
+import { Action } from "../../Shared/Telemetry/TelemetryConstants";
+import * as TelemetryProcessor from "../../Shared/Telemetry/TelemetryProcessor";
+import DocumentId from "../Tree/DocumentId";
+import ObjectId from "../Tree/ObjectId";
+import DocumentsTab from "./DocumentsTab";
 
 export default class MongoDocumentsTab extends DocumentsTab {
   public collection: ViewModels.Collection;

src/Explorer/Tree/Collection.ts

@@ -513,9 +513,7 @@ export default class Collection implements ViewModels.Collection {
         tabKind: ViewModels.CollectionTabKind.Documents,
         title: "Documents",
         tabPath: "",
-
         collection: this,
-
         node: this,
         hashLocation: `${Constants.HashRoutePrefixes.collectionsWithIds(this.databaseId, this.id())}/mongoDocuments`,
         isActive: ko.observable(false),

src/Explorer/Tree/Database.ts

@@ -172,6 +172,27 @@ export default class Database implements ViewModels.Database {
   public async loadCollections(): Promise<void> {
     const collectionVMs: Collection[] = [];
     const collections: DataModels.Collection[] = await readCollections(this.id());
+    // TODO Remove
+    // This is a hack to make Mongo collections read via ARM have a SQL-ish partitionKey property
+    if (userContext.apiType === "Mongo") {
+      collections.map((collection) => {
+        if (collection.shardKey) {
+          const shardKey = Object.keys(collection.shardKey)[0];
+          collection.partitionKey = {
+            version: undefined,
+            kind: "Hash",
+            paths: [`/"$v"/"${shardKey.split(".").join(`"/"$v"/"`)}"/"$v"`],
+          };
+        } else {
+          collection.partitionKey = {
+            paths: ["/'$v'/'_partitionKey'/'$v'"],
+            kind: "Hash",
+            version: 2,
+            systemKey: true,
+          };
+        }
+      });
+    }
     const deltaCollections = this.getDeltaCollections(collections);
 
     collections.forEach((collection: DataModels.Collection) => {