From 7e0c4b72909110c9862b6ee76fec526f08247d83 Mon Sep 17 00:00:00 2001
From: t-tarabhatia <84812001+t-tarabhatia@users.noreply.github.com>
Date: Thu, 29 Jul 2021 06:48:03 -0700
Subject: [PATCH] Add changes to Partition Key A/B Test (#954)

---
 src/Common/Constants.ts                   |  1 +
 src/Explorer/Panes/AddCollectionPanel.tsx | 19 ++++++++++++++-----
 src/Platform/Hosted/extractFeatures.ts    |  2 ++
 src/hooks/useKnockoutExplorer.ts          |  6 ++++++
 4 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/src/Common/Constants.ts b/src/Common/Constants.ts
index fc97d559d..cbdd2da8e 100644
--- a/src/Common/Constants.ts
+++ b/src/Common/Constants.ts
@@ -95,6 +95,7 @@ export class Flights {
   public static readonly MongoIndexing = "mongoindexing";
   public static readonly AutoscaleTest = "autoscaletest";
   public static readonly PartitionKeyTest = "partitionkeytest";
+  public static readonly PKPartitionKeyTest = "pkpartitionkeytest";
 }
 
 export class AfecFeatures {
diff --git a/src/Explorer/Panes/AddCollectionPanel.tsx b/src/Explorer/Panes/AddCollectionPanel.tsx
index f0ce4f4d1..39d7ad2c8 100644
--- a/src/Explorer/Panes/AddCollectionPanel.tsx
+++ b/src/Explorer/Panes/AddCollectionPanel.tsx
@@ -113,11 +113,7 @@ export class AddCollectionPanel extends React.Component<AddCollectionPanelProps,
       collectionId: "",
       enableIndexing: true,
       isSharded: userContext.apiType !== "Tables",
-      partitionKey:
-        (userContext.features.partitionKeyDefault && userContext.apiType === "SQL") ||
-        (userContext.features.partitionKeyDefault && userContext.apiType === "Mongo")
-          ? "/id"
-          : "",
+      partitionKey: this.getPartitionKey(),
       enableDedicatedThroughput: false,
       createMongoWildCardIndex: isCapabilityEnabled("EnableMongo"),
       useHashV2: false,
@@ -815,6 +811,19 @@ export class AddCollectionPanel extends React.Component<AddCollectionPanelProps,
     return tooltipText;
   }
 
+  private getPartitionKey(): string {
+    if (userContext.apiType !== "SQL" && userContext.apiType !== "Mongo") {
+      return "";
+    }
+    if (userContext.features.partitionKeyDefault) {
+      return userContext.apiType === "SQL" ? "/id" : "_id";
+    }
+    if (userContext.features.partitionKeyDefault2) {
+      return userContext.apiType === "SQL" ? "/pk" : "pk";
+    }
+    return "";
+  }
+
   private getPartitionKeySubtext(): string {
     if (
       userContext.features.partitionKeyDefault &&
diff --git a/src/Platform/Hosted/extractFeatures.ts b/src/Platform/Hosted/extractFeatures.ts
index f1f0115a5..4a85dca8d 100644
--- a/src/Platform/Hosted/extractFeatures.ts
+++ b/src/Platform/Hosted/extractFeatures.ts
@@ -10,6 +10,7 @@ export type Features = {
   readonly enableSchema: boolean;
   autoscaleDefault: boolean;
   partitionKeyDefault: boolean;
+  partitionKeyDefault2: boolean;
   readonly enableSDKoperations: boolean;
   readonly enableSpark: boolean;
   readonly enableTtl: boolean;
@@ -72,5 +73,6 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     ttl90Days: "true" === get("ttl90days"),
     autoscaleDefault: "true" === get("autoscaledefault"),
     partitionKeyDefault: "true" === get("partitionkeytest"),
+    partitionKeyDefault2: "true" === get("pkpartitionkeytest"),
   };
 }
diff --git a/src/hooks/useKnockoutExplorer.ts b/src/hooks/useKnockoutExplorer.ts
index 86200808c..686391376 100644
--- a/src/hooks/useKnockoutExplorer.ts
+++ b/src/hooks/useKnockoutExplorer.ts
@@ -331,6 +331,12 @@ function updateContextsFromPortalMessage(inputs: DataExplorerInputsFrame) {
     if (inputs.flights.indexOf(Flights.PartitionKeyTest) !== -1) {
       userContext.features.partitionKeyDefault = true;
     }
+    if (inputs.flights.indexOf(Flights.PartitionKeyTest) !== -1) {
+      userContext.features.partitionKeyDefault = true;
+    }
+    if (inputs.flights.indexOf(Flights.PKPartitionKeyTest) !== -1) {
+      userContext.features.partitionKeyDefault2 = true;
+    }
   }
 }