From 7bdc31aa675af54506d3e1aaf0a34f36539ab011 Mon Sep 17 00:00:00 2001
From: Steve Faulkner <southpolesteve@gmail.com>
Date: Tue, 13 Apr 2021 12:03:25 -0500
Subject: [PATCH] Fix for Mongo shard key loaded via ARM (#657)

---
 src/Contracts/DataModels.ts            |  4 ++++
 src/Explorer/Tabs/MongoDocumentsTab.ts | 25 ++++++++++++-------------
 src/Explorer/Tree/Collection.ts        |  2 --
 src/Explorer/Tree/Database.ts          | 21 +++++++++++++++++++++
 4 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/src/Contracts/DataModels.ts b/src/Contracts/DataModels.ts
index 4e23a64dd..4e3085b45 100644
--- a/src/Contracts/DataModels.ts
+++ b/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;
diff --git a/src/Explorer/Tabs/MongoDocumentsTab.ts b/src/Explorer/Tabs/MongoDocumentsTab.ts
index 61fce5059..f5b2f2959 100644
--- a/src/Explorer/Tabs/MongoDocumentsTab.ts
+++ b/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;
diff --git a/src/Explorer/Tree/Collection.ts b/src/Explorer/Tree/Collection.ts
index 83fdc47d1..49945b22c 100644
--- a/src/Explorer/Tree/Collection.ts
+++ b/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),
diff --git a/src/Explorer/Tree/Database.ts b/src/Explorer/Tree/Database.ts
index 91df402b4..e357fbca8 100644
--- a/src/Explorer/Tree/Database.ts
+++ b/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) => {