From 8200cc521f68088dc28bc0ef24e657cf56f96731 Mon Sep 17 00:00:00 2001
From: Laurent Nguyen <laurent.nguyen@microsoft.com>
Date: Fri, 26 Jun 2020 16:52:54 +0200
Subject: [PATCH] Switch to Graph explorer gremlin queries to use id and pk
 inside single quoted strings (#57)

---
 .../GraphExplorer.test.tsx                    | 81 ++++++++++++++++---
 .../GraphExplorerComponent/GraphExplorer.tsx  | 37 ++++++---
 2 files changed, 100 insertions(+), 18 deletions(-)

diff --git a/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.test.tsx b/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.test.tsx
index b147ab10c..56c3dc956 100644
--- a/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.test.tsx
+++ b/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.test.tsx
@@ -11,7 +11,6 @@ import * as ViewModels from "../../../Contracts/ViewModels";
 import * as DataModels from "../../../Contracts/DataModels";
 import * as StorageUtility from "../../../Shared/StorageUtility";
 import GraphTab from "../../Tabs/GraphTab";
-import DocumentClientUtilityBase from "../../../Common/DocumentClientUtilityBase";
 import { ConsoleDataType } from "../../Menus/NotificationConsole/NotificationConsoleComponent";
 
 describe("Check whether query result is vertex array", () => {
@@ -59,10 +58,74 @@ describe("Check whether query result is edge-vertex array", () => {
 
 describe("Create proper pkid pair", () => {
   it("should enclose string pk with quotes", () => {
-    expect(GraphExplorer.generatePkIdPair("test", "id")).toEqual('["test", "id"]');
+    expect(GraphExplorer.generatePkIdPair("test", "id")).toEqual("['test', 'id']");
   });
+
   it("should not enclose non-string pk with quotes", () => {
-    expect(GraphExplorer.generatePkIdPair(2, "id")).toEqual('[2, "id"]');
+    expect(GraphExplorer.generatePkIdPair(2, "id")).toEqual("[2, 'id']");
+  });
+});
+
+describe("getPkIdFromDocumentId", () => {
+  const createFakeDoc = (override: any) => ({
+    _rid: "_rid",
+    _self: "_self",
+    _etag: "_etag",
+    _ts: 1234,
+    ...override
+  });
+
+  it("should create pkid pair from non-partitioned graph", () => {
+    const doc = createFakeDoc({ id: "id" });
+    expect(GraphExplorer.getPkIdFromDocumentId(doc, undefined)).toEqual("'id'");
+    expect(GraphExplorer.getPkIdFromDocumentId(doc, "_partitiongKey")).toEqual("'id'");
+  });
+
+  it("should create pkid pair from partitioned graph (pk as string)", () => {
+    const doc = createFakeDoc({ id: "id", mypk: "pkvalue" });
+    expect(GraphExplorer.getPkIdFromDocumentId(doc, "mypk")).toEqual("['pkvalue', 'id']");
+  });
+
+  it("should create pkid pair from partitioned graph (pk as valid array value)", () => {
+    const doc = createFakeDoc({ id: "id", mypk: [{ id: "someid", _value: "pkvalue" }] });
+    expect(GraphExplorer.getPkIdFromDocumentId(doc, "mypk")).toEqual("['pkvalue', 'id']");
+  });
+
+  it("should error if id is not a string", () => {
+    const doc = createFakeDoc({ id: { foo: 1 } });
+    try {
+      GraphExplorer.getPkIdFromDocumentId(doc, undefined);
+      expect(true).toBe(false);
+    } catch (e) {
+      expect(true).toBe(true);
+    }
+  });
+
+  it("should error if pk not string nor non-empty array", () => {
+    let doc = createFakeDoc({ mypk: { foo: 1 } });
+
+    try {
+      GraphExplorer.getPkIdFromDocumentId(doc, "mypk");
+    } catch (e) {
+      expect(true).toBe(true);
+    }
+
+    doc = createFakeDoc({ mypk: [] });
+    try {
+      GraphExplorer.getPkIdFromDocumentId(doc, "mypk");
+      expect(true).toBe(false);
+    } catch (e) {
+      expect(true).toBe(true);
+    }
+
+    // Array must be [{ id: string, _value: string }]
+    doc = createFakeDoc({ mypk: [{ foo: 1 }] });
+    try {
+      GraphExplorer.getPkIdFromDocumentId(doc, "mypk");
+      expect(true).toBe(false);
+    } catch (e) {
+      expect(true).toBe(true);
+    }
   });
 });
 
@@ -253,11 +316,11 @@ describe("GraphExplorer", () => {
     };
 
     const createFetchOutEQuery = (vertexId: string, limit: number): string => {
-      return `g.V("${vertexId}").outE().limit(${limit}).as('e').inV().as('v').select('e', 'v')`;
+      return `g.V('${vertexId}').outE().limit(${limit}).as('e').inV().as('v').select('e', 'v')`;
     };
 
     const createFetchInEQuery = (vertexId: string, limit: number): string => {
-      return `g.V("${vertexId}").inE().limit(${limit}).as('e').outV().as('v').select('e', 'v')`;
+      return `g.V('${vertexId}').inE().limit(${limit}).as('e').outV().as('v').select('e', 'v')`;
     };
 
     const isVisible = (selector: string): boolean => {
@@ -293,7 +356,7 @@ describe("GraphExplorer", () => {
     describe("Load Graph button", () => {
       beforeEach(async done => {
         const backendResponses: BackendResponses = {};
-        backendResponses["g.V()"] = backendResponses['g.V("1")'] = {
+        backendResponses["g.V()"] = backendResponses["g.V('1')"] = {
           response: [{ id: "1", type: "vertex" }],
           isLast: false
         };
@@ -341,7 +404,7 @@ describe("GraphExplorer", () => {
     describe("Execute Gremlin Query button", () => {
       beforeEach(done => {
         const backendResponses: BackendResponses = {};
-        backendResponses["g.V()"] = backendResponses['g.V("2")'] = {
+        backendResponses["g.V()"] = backendResponses["g.V('2')"] = {
           response: [{ id: "2", type: "vertex" }],
           isLast: false
         };
@@ -411,7 +474,7 @@ describe("GraphExplorer", () => {
       beforeEach(done => {
         const backendResponses: BackendResponses = {};
         // TODO Make this less dependent on spaces, order and quotes
-        backendResponses["g.V()"] = backendResponses[`g.V("${node1Id}","${node2Id}")`] = {
+        backendResponses["g.V()"] = backendResponses[`g.V('${node1Id}','${node2Id}')`] = {
           response: [
             {
               id: node1Id,
@@ -667,7 +730,7 @@ describe("GraphExplorer", () => {
     describe("when isGraphAutoVizDisabled setting is true (autoviz disabled)", () => {
       beforeEach(done => {
         const backendResponses: BackendResponses = {};
-        backendResponses["g.V()"] = backendResponses['g.V("3")'] = {
+        backendResponses["g.V()"] = backendResponses["g.V('3')"] = {
           response: [{ id: "3", type: "vertex" }],
           isLast: true
         };
diff --git a/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.tsx b/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.tsx
index 7e5bf326e..500672c0f 100644
--- a/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.tsx
+++ b/src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.tsx
@@ -327,8 +327,8 @@ export class GraphExplorer extends React.Component<GraphExplorerProps, GraphExpl
    * @param id
    */
   public static generatePkIdPair(pk: PartitionKeyValueType, id: string) {
-    const pkStr = typeof pk === "string" ? `"${pk}"` : `${pk}`;
-    return `[${pkStr}, "${GraphUtil.escapeDoubleQuotes(id)}"]`;
+    const pkStr = typeof pk === "string" ? `'${pk}'` : `${pk}`;
+    return `[${pkStr}, '${GraphUtil.escapeSingleQuotes(id)}']`;
   }
 
   public updateVertexProperties(editedProperties: EditedProperties): Q.Promise<GremlinClient.GremlinRequestResult> {
@@ -1335,7 +1335,7 @@ export class GraphExplorer extends React.Component<GraphExplorerProps, GraphExpl
       const pk = v.properties[this.props.collectionPartitionKeyProperty][0].value;
       return GraphExplorer.generatePkIdPair(pk, v.id);
     } else {
-      return `"${GraphUtil.escapeDoubleQuotes(v.id)}"`;
+      return `'${GraphUtil.escapeSingleQuotes(v.id)}'`;
     }
   }
 
@@ -1361,15 +1361,34 @@ export class GraphExplorer extends React.Component<GraphExplorerProps, GraphExpl
   /**
    * If collection is not partitioned, return 'id'.
    * If collection is partitioned, return pk-id pair.
+   * public for testing purposes
    * @param vertex
    * @return id
    */
-  private getPkIdFromDocumentId(d: DataModels.DocumentId): string {
-    if (this.props.collectionPartitionKeyProperty && d.hasOwnProperty(this.props.collectionPartitionKeyProperty)) {
-      const pk = (d as any)[this.props.collectionPartitionKeyProperty];
-      return GraphExplorer.generatePkIdPair(pk, d.id);
+  public static getPkIdFromDocumentId(d: DataModels.DocumentId, collectionPartitionKeyProperty: string): string {
+    let { id } = d;
+    if (typeof id !== "string") {
+      const error = `Vertex id is not a string: ${JSON.stringify(id)}.`;
+      NotificationConsoleUtils.logConsoleMessage(ConsoleDataType.Error, error);
+      throw new Error(error);
+    }
+
+    if (collectionPartitionKeyProperty && d.hasOwnProperty(collectionPartitionKeyProperty)) {
+      let pk = (d as any)[collectionPartitionKeyProperty];
+      if (typeof pk !== "string") {
+        if (Array.isArray(pk) && pk.length > 0) {
+          // pk is [{ id: 'id', _value: 'value' }]
+          pk = pk[0]["_value"];
+        } else {
+          const error = `Vertex pk is not a string nor a non-empty array: ${JSON.stringify(pk)}.`;
+          NotificationConsoleUtils.logConsoleMessage(ConsoleDataType.Error, error);
+          throw new Error(error);
+        }
+      }
+
+      return GraphExplorer.generatePkIdPair(pk, id);
     } else {
-      return `"${GraphUtil.escapeDoubleQuotes(d.id)}"`;
+      return `'${GraphUtil.escapeSingleQuotes(id)}'`;
     }
   }
 
@@ -1769,7 +1788,7 @@ export class GraphExplorer extends React.Component<GraphExplorerProps, GraphExpl
         const documents = results.documents || [];
         return documents.map(
           (item: DataModels.DocumentId) => {
-            return this.getPkIdFromDocumentId(item);
+            return GraphExplorer.getPkIdFromDocumentId(item, this.props.collectionPartitionKeyProperty);
           },
           (reason: any) => {
             // Failure