Display stored procedures in graphs (#639)

* Display stored procedures in graphs

* Pass client instead of collectionTreeItem

src/docdb/tree/DocDBStoredProcedureTreeItem.ts

@@ -8,7 +8,6 @@ import * as vscode from "vscode";
 import { IAzureTreeItem, IAzureNode, UserCancelledError, DialogResponses } from 'vscode-azureextensionui';
 import { ProcedureMeta, DocumentClient } from 'documentdb';
 import { getDocumentClient } from '../getDocumentClient';
-import { DocDBCollectionTreeItem } from './DocDBCollectionTreeItem';
 
 /**
  * Represents a Cosmos DB DocumentDB (SQL) stored procedure
@@ -18,7 +17,7 @@ export class DocDBStoredProcedureTreeItem implements IAzureTreeItem {
     public readonly contextValue: string = DocDBStoredProcedureTreeItem.contextValue;
     public readonly commandId: string = 'cosmosDB.openStoredProcedure';
 
-    constructor(private _endpoint: string, private _masterKey: string, private _isEmulator: boolean, private _collection: DocDBCollectionTreeItem, public procedure: ProcedureMeta) {
+    constructor(private _endpoint: string, private _masterKey: string, private _isEmulator: boolean, private _client: DocumentClient, public procedure: ProcedureMeta) {
     }
 
     public get id(): string {
@@ -34,8 +33,7 @@ export class DocDBStoredProcedureTreeItem implements IAzureTreeItem {
     }
 
     public async update(newProcBody: string): Promise<string> {
-        const client: DocumentClient = this._collection.getDocumentClient();
-        this.procedure = await new Promise<ProcedureMeta>((resolve, reject) => client.replaceStoredProcedure(
+        this.procedure = await new Promise<ProcedureMeta>((resolve, reject) => this._client.replaceStoredProcedure(
             this.link,
             { body: newProcBody, id: this.procedure.id },
             (err, updated: ProcedureMeta) => {

src/docdb/tree/DocDBStoredProceduresTreeItem.ts

@@ -11,6 +11,7 @@ import { IAzureTreeItem, UserCancelledError, IAzureNode } from 'vscode-azureexte
 import { DocDBStoredProcedureTreeItem } from './DocDBStoredProcedureTreeItem';
 import { defaultStoredProcedure } from '../../constants';
 import { DocDBCollectionTreeItem } from './DocDBCollectionTreeItem';
+import { GraphCollectionTreeItem } from '../../graph/tree/GraphCollectionTreeItem';
 
 /**
  * This class represents the DocumentDB "Stored Procedures" node in the tree
@@ -20,12 +21,12 @@ export class DocDBStoredProceduresTreeItem extends DocDBTreeItemBase<ProcedureMe
     public readonly contextValue: string = DocDBStoredProceduresTreeItem.contextValue;
     public readonly childTypeLabel: string = "Stored Procedure";
 
-    constructor(endpoint: string, masterKey: string, private _collection: DocDBCollectionTreeItem, isEmulator: boolean) {
+    constructor(endpoint: string, masterKey: string, private _collection: DocDBCollectionTreeItem | GraphCollectionTreeItem, isEmulator: boolean) {
         super(endpoint, masterKey, isEmulator);
     }
 
     public initChild(resource: ProcedureMeta): IAzureTreeItem {
-        return new DocDBStoredProcedureTreeItem(this.documentEndpoint, this.masterKey, this.isEmulator, this._collection, resource);
+        return new DocDBStoredProcedureTreeItem(this.documentEndpoint, this.masterKey, this.isEmulator, this._collection.getDocumentClient(), resource);
     }
 
     public get iconPath(): string | vscode.Uri | { light: string | vscode.Uri; dark: string | vscode.Uri } {

src/graph/registerGraphCommands.ts

@@ -9,6 +9,7 @@ import { GraphAccountTreeItem } from "./tree/GraphAccountTreeItem";
 import { GraphDatabaseTreeItem } from "./tree/GraphDatabaseTreeItem";
 import { GraphCollectionTreeItem } from "./tree/GraphCollectionTreeItem";
 import { GraphViewsManager } from "./GraphViewsManager";
+import { GraphTreeItem } from "./tree/GraphTreeItem";
 
 export function registerGraphCommands(context: vscode.ExtensionContext, actionHandler: AzureActionHandler, tree: AzureTreeDataProvider): void {
     let graphViewsManager = new GraphViewsManager(context);
@@ -37,9 +38,9 @@ export function registerGraphCommands(context: vscode.ExtensionContext, actionHa
         }
         await node.deleteNode();
     });
-    actionHandler.registerCommand('cosmosDB.openGraphExplorer', async (node: IAzureNode<GraphCollectionTreeItem>) => {
+    actionHandler.registerCommand('cosmosDB.openGraphExplorer', async (node: IAzureNode<GraphTreeItem>) => {
         if (!node) {
-            node = <IAzureNode<GraphCollectionTreeItem>>await tree.showNodePicker(GraphCollectionTreeItem.contextValue);
+            node = <IAzureNode<GraphTreeItem>>await tree.showNodePicker(GraphCollectionTreeItem.contextValue);
         }
         await node.treeItem.showExplorer(graphViewsManager);
     });

src/graph/tree/GraphCollectionTreeItem.ts

@@ -5,23 +5,28 @@
 
 import * as path from 'path';
 import { IAzureTreeItem, IAzureNode, UserCancelledError, DialogResponses } from 'vscode-azureextensionui';
-import { GraphViewsManager } from '../GraphViewsManager';
-import { GraphConfiguration } from '../GraphConfiguration';
 import * as vscode from 'vscode';
-import { CollectionMeta } from 'documentdb';
+import { CollectionMeta, DocumentClient } from 'documentdb';
 import { GraphDatabaseTreeItem } from './GraphDatabaseTreeItem';
+import { GraphTreeItem } from './GraphTreeItem';
+import { DocDBStoredProceduresTreeItem } from '../../docdb/tree/DocDBStoredProceduresTreeItem';
+import { getDocumentClient } from '../../docdb/getDocumentClient';
 
 export class GraphCollectionTreeItem implements IAzureTreeItem {
     public static contextValue: string = "cosmosDBGraph";
     public readonly contextValue: string = GraphCollectionTreeItem.contextValue;
-    public readonly commandId: string = 'cosmosDB.openGraphExplorer';
+
+    private readonly _graphTreeItem: GraphTreeItem;
+    private readonly _storedProceduresTreeItem: DocDBStoredProceduresTreeItem;
 
     private readonly _database: GraphDatabaseTreeItem;
     private readonly _collection: CollectionMeta;
 
-    constructor(database: GraphDatabaseTreeItem, collection: CollectionMeta) {
+    constructor(database: GraphDatabaseTreeItem, collection: CollectionMeta, private _documentEndpoint: string, private _masterKey: string, private _isEmulator: boolean) {
         this._database = database;
         this._collection = collection;
+        this._graphTreeItem = new GraphTreeItem(this._database, this._collection);
+        this._storedProceduresTreeItem = new DocDBStoredProceduresTreeItem(this._documentEndpoint, this._masterKey, this, this._isEmulator);
     }
 
     public get id(): string {
@@ -43,6 +48,18 @@ export class GraphCollectionTreeItem implements IAzureTreeItem {
         };
     }
 
+    public getDocumentClient(): DocumentClient {
+        return getDocumentClient(this._documentEndpoint, this._masterKey, this._isEmulator);
+    }
+
+    public async loadMoreChildren(_node: IAzureNode<IAzureTreeItem>, _clearCache: boolean): Promise<IAzureTreeItem[]> {
+        return [this._graphTreeItem, this._storedProceduresTreeItem];
+    }
+
+    public hasMoreChildren(): boolean {
+        return false;
+    }
+
     public async deleteTreeItem(_node: IAzureNode): Promise<void> {
         const message: string = `Are you sure you want to delete graph '${this.label}' and its contents?`;
         const result = await vscode.window.showWarningMessage(message, { modal: true }, DialogResponses.deleteResponse, DialogResponses.cancel);
@@ -57,15 +74,4 @@ export class GraphCollectionTreeItem implements IAzureTreeItem {
             throw new UserCancelledError();
         }
     }
-
-    public async showExplorer(graphViewsManager: GraphViewsManager): Promise<void> {
-        await graphViewsManager.showGraphViewer(this.label, <GraphConfiguration>{
-            documentEndpoint: this._database.documentEndpoint,
-            gremlinEndpoint: this._database.gremlinEndpoint,
-            possibleGremlinEndpoints: this._database.possibleGremlinEndpoints,
-            databaseName: this._database.label,
-            graphName: this.label,
-            key: this._database.masterKey
-        });
-    }
 }

src/graph/tree/GraphDatabaseTreeItem.ts

@@ -19,7 +19,7 @@ export class GraphDatabaseTreeItem extends DocDBDatabaseTreeItemBase {
     }
 
     public initChild(collection: CollectionMeta): IAzureTreeItem {
-        return new GraphCollectionTreeItem(this, collection);
+        return new GraphCollectionTreeItem(this, collection, this.documentEndpoint, this.masterKey, this.isEmulator);
     }
 
     // Gremlin endpoint, if definitely known

src/graph/tree/GraphTreeItem.ts

@@ -0,0 +1,56 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as path from 'path';
+import { IAzureTreeItem } from 'vscode-azureextensionui';
+import { GraphViewsManager } from '../GraphViewsManager';
+import { GraphConfiguration } from '../GraphConfiguration';
+import * as vscode from 'vscode';
+import { CollectionMeta } from 'documentdb';
+import { GraphDatabaseTreeItem } from './GraphDatabaseTreeItem';
+
+export class GraphTreeItem implements IAzureTreeItem {
+    public static contextValue: string = "cosmosDBGraphGraph";
+    public readonly contextValue: string = GraphTreeItem.contextValue;
+    public readonly commandId: string = 'cosmosDB.openGraphExplorer';
+
+    private readonly _database: GraphDatabaseTreeItem;
+    private readonly _collection: CollectionMeta;
+
+    constructor(database: GraphDatabaseTreeItem, collection: CollectionMeta) {
+        this._database = database;
+        this._collection = collection;
+    }
+
+    public get id(): string {
+        return this._collection.id;
+    }
+
+    public get label(): string {
+        return "Graph";
+    }
+
+    public get link(): string {
+        return this._collection._self;
+    }
+
+    public get iconPath(): string | vscode.Uri | { light: string | vscode.Uri; dark: string | vscode.Uri } {
+        return {
+            light: path.join(__filename, '..', '..', '..', '..', '..', 'resources', 'icons', 'theme-agnostic', 'Collection.svg'),
+            dark: path.join(__filename, '..', '..', '..', '..', '..', 'resources', 'icons', 'theme-agnostic', 'Collection.svg'),
+        };
+    }
+
+    public async showExplorer(graphViewsManager: GraphViewsManager): Promise<void> {
+        await graphViewsManager.showGraphViewer(this.label, <GraphConfiguration>{
+            documentEndpoint: this._database.documentEndpoint,
+            gremlinEndpoint: this._database.gremlinEndpoint,
+            possibleGremlinEndpoints: this._database.possibleGremlinEndpoints,
+            databaseName: this._database.label,
+            graphName: this.label,
+            key: this._database.masterKey
+        });
+    }
+}