Support data plane RBAC for Gremlin API (#2182)

* Refactor logic for determining if we should use data plane RBAC to a
common function.

* Support RBAC for gremlin API.

* Refactor to use common function.

* Fix unit tests.

* Move test function inside test scope.

* Minor clean ups.

* Reinstate utf8ToB64 function in case this breaks a corner case.

src/Explorer/DataSamples/ContainerSampleGenerator.ts

@@ -1,3 +1,4 @@
+import { useDataplaneRbacAuthorization } from "Utils/AuthorizationUtils";
 import { createCollection } from "../../Common/dataAccess/createCollection";
 import { createDocument } from "../../Common/dataAccess/createDocument";
 import { createDocument as createMongoDocument } from "../../Common/MongoProxyClient";
@@ -90,12 +91,13 @@ export class ContainerSampleGenerator {
       }
       const { databaseAccount: account } = userContext;
       const databaseId = collection.databaseId;
+
       const gremlinClient = new GremlinClient();
       gremlinClient.initialize({
         endpoint: `wss://${GraphTab.getGremlinEndpoint(account)}`,
         databaseId: databaseId,
         collectionId: collection.id(),
-        masterKey: userContext.masterKey || "",
+        password: useDataplaneRbacAuthorization(userContext) ? userContext.aadToken : userContext.masterKey || "",
         maxResultSize: 100,
       });
 

src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.test.tsx

@@ -163,8 +163,7 @@ describe("GraphExplorer", () => {
       graphBackendEndpoint: "graphBackendEndpoint",
       databaseId: "databaseId",
       collectionId: "collectionId",
-      masterKey: "masterKey",
-
+      password: "password",
       onLoadStartKey: 0,
       onLoadStartKeyChange: (newKey: number): void => {},
       resourceId: "resourceId",

src/Explorer/Graph/GraphExplorerComponent/GraphExplorer.tsx

@@ -59,7 +59,7 @@ export interface GraphExplorerProps {
   graphBackendEndpoint: string;
   databaseId: string;
   collectionId: string;
-  masterKey: string;
+  password: string;
 
   onLoadStartKey: number;
   onLoadStartKeyChange: (newKey: number) => void;
@@ -1300,7 +1300,7 @@ export class GraphExplorer extends React.Component<GraphExplorerProps, GraphExpl
       endpoint: `wss://${this.props.graphBackendEndpoint}`,
       databaseId: this.props.databaseId,
       collectionId: this.props.collectionId,
-      masterKey: this.props.masterKey,
+      password: this.props.password,
       maxResultSize: GraphExplorer.MAX_RESULT_SIZE,
     });
   }

src/Explorer/Graph/GraphExplorerComponent/GremlinClient.test.ts

@@ -8,28 +8,28 @@ describe("Gremlin Client", () => {
     endpoint: null,
     collectionId: null,
     databaseId: null,
-    masterKey: null,
     maxResultSize: 10000,
+    password: null,
   };
 
-  it("should use databaseId, collectionId and masterKey to authenticate", () => {
+  it("should use databaseId, collectionId and password to authenticate", () => {
     const collectionId = "collectionId";
     const databaseId = "databaseId";
-    const masterKey = "masterKey";
+    const testPassword = "password";
     const gremlinClient = new GremlinClient();
 
     gremlinClient.initialize({
       endpoint: null,
       collectionId,
       databaseId,
-      masterKey,
       maxResultSize: 0,
+      password: testPassword,
     });
 
     // User must includes these values
     expect(gremlinClient.client.params.user.indexOf(collectionId)).not.toBe(-1);
     expect(gremlinClient.client.params.user.indexOf(databaseId)).not.toBe(-1);
-    expect(gremlinClient.client.params.password).toEqual(masterKey);
+    expect(gremlinClient.client.params.password).toEqual(testPassword);
   });
 
   it("should aggregate RU charges across multiple responses", (done) => {

src/Explorer/Graph/GraphExplorerComponent/GremlinClient.ts

@@ -11,8 +11,8 @@ export interface GremlinClientParameters {
   endpoint: string;
   databaseId: string;
   collectionId: string;
-  masterKey: string;
   maxResultSize: number;
+  password: string;
 }
 
 export interface GremlinRequestResult {
@@ -43,7 +43,7 @@ export class GremlinClient {
     this.client = new GremlinSimpleClient({
       endpoint: params.endpoint,
       user: `/dbs/${params.databaseId}/colls/${params.collectionId}`,
-      password: params.masterKey,
+      password: params.password,
       successCallback: (result: Result) => {
         this.storePendingResult(result);
         this.flushResult(result.requestId);

src/Explorer/Graph/GraphExplorerComponent/GremlinSimpleClient.test.ts

@@ -5,11 +5,11 @@
 
 import * as sinon from "sinon";
 import {
+  GremlinRequestMessage,
+  GremlinResponseMessage,
   GremlinSimpleClient,
   GremlinSimpleClientParameters,
   Result,
-  GremlinRequestMessage,
-  GremlinResponseMessage,
 } from "./GremlinSimpleClient";
 
 describe("Gremlin Simple Client", () => {

src/Explorer/Tabs/GraphTab.tsx

@@ -45,7 +45,7 @@ export interface IGraphConfig {
 
 interface GraphTabOptions extends ViewModels.TabOptions {
   account: DatabaseAccount;
-  masterKey: string;
+  password: string;
   collectionId: string;
   databaseId: string;
   collectionPartitionKeyProperty: string;
@@ -107,7 +107,7 @@ export default class GraphTab extends TabsBase {
       graphBackendEndpoint: GraphTab.getGremlinEndpoint(options.account),
       databaseId: options.databaseId,
       collectionId: options.collectionId,
-      masterKey: options.masterKey,
+      password: options.password,
       onLoadStartKey: options.onLoadStartKey,
       onLoadStartKeyChange: (onLoadStartKey: number): void => {
         if (onLoadStartKey === undefined) {

src/Explorer/Tree/Collection.ts

@@ -8,6 +8,7 @@ import {
 import { useNotebook } from "Explorer/Notebook/useNotebook";
 import { DocumentsTabV2 } from "Explorer/Tabs/DocumentsTabV2/DocumentsTabV2";
 import { isFabricMirrored } from "Platform/Fabric/FabricUtil";
+import { useDataplaneRbacAuthorization } from "Utils/AuthorizationUtils";
 import * as ko from "knockout";
 import * as _ from "underscore";
 import * as Constants from "../../Common/Constants";
@@ -479,9 +480,8 @@ export default class Collection implements ViewModels.Collection {
         node: this,
         title: title,
         tabPath: "",
-
+        password: useDataplaneRbacAuthorization(userContext) ? userContext.aadToken : userContext.masterKey || "",
         collection: this,
-        masterKey: userContext.masterKey || "",
         collectionPartitionKeyProperty: this.partitionKeyProperties?.[0],
         collectionId: this.id(),
         databaseId: this.databaseId,
@@ -737,7 +737,7 @@ export default class Collection implements ViewModels.Collection {
       title: title,
       tabPath: "",
       collection: this,
-      masterKey: userContext.masterKey || "",
+      password: useDataplaneRbacAuthorization(userContext) ? userContext.aadToken : userContext.masterKey || "",
       collectionPartitionKeyProperty: this.partitionKeyProperties?.[0],
       collectionId: this.id(),
       databaseId: this.databaseId,

src/Utils/APITypeUtils.ts

@@ -91,5 +91,5 @@ export const getItemName = (): string => {
 };
 
 export const isDataplaneRbacSupported = (apiType: string): boolean => {
-  return apiType === "SQL" || apiType === "Tables";
+  return apiType === "SQL" || apiType === "Tables" || apiType === "Gremlin";
 };

src/Utils/AuthorizationUtils.test.ts

@@ -104,7 +104,7 @@ describe("AuthorizationUtils", () => {
 
     it("should return true if dataPlaneRbacEnabled is set to true and API supports RBAC", () => {
       setAadDataPlane(false);
-      ["SQL", "Tables"].forEach((type) => {
+      ["SQL", "Tables", "Gremlin"].forEach((type) => {
         updateUserContext({
           dataPlaneRbacEnabled: true,
           apiType: type as ApiType,
@@ -115,7 +115,7 @@ describe("AuthorizationUtils", () => {
 
     it("should return false if dataPlaneRbacEnabled is set to true and API does not support RBAC", () => {
       setAadDataPlane(false);
-      ["Mongo", "Gremlin", "Cassandra", "Postgres", "VCoreMongo"].forEach((type) => {
+      ["Mongo", "Cassandra", "Postgres", "VCoreMongo"].forEach((type) => {
         updateUserContext({
           dataPlaneRbacEnabled: true,
           apiType: type as ApiType,