From cb1d60cc90438b6af2298f518b7b01bfa6aea9b3 Mon Sep 17 00:00:00 2001
From: Tanuj Mittal <tamitta@microsoft.com>
Date: Wed, 26 May 2021 13:01:13 -0700
Subject: [PATCH] Hide hosted shells and schema analyzer if VNET, Firewall or
 Private endpoints is enabled (#826)

* Disable Schema Analyzer if VNET or Firewall is enabled

* Add support for private endpoint connections

* Fix lint warning
---
 src/Common/DatabaseAccountUtility.ts      | 17 +++++++++++++++++
 src/Contracts/DataModels.ts               |  1 +
 src/Explorer/Explorer.tsx                 |  8 +++-----
 src/Explorer/Tree/ResourceTreeAdapter.tsx |  7 ++-----
 4 files changed, 23 insertions(+), 10 deletions(-)
 create mode 100644 src/Common/DatabaseAccountUtility.ts

diff --git a/src/Common/DatabaseAccountUtility.ts b/src/Common/DatabaseAccountUtility.ts
new file mode 100644
index 000000000..693f581f8
--- /dev/null
+++ b/src/Common/DatabaseAccountUtility.ts
@@ -0,0 +1,17 @@
+import { userContext } from "../UserContext";
+
+function isVirtualNetworkFilterEnabled() {
+  return userContext.databaseAccount?.properties?.isVirtualNetworkFilterEnabled;
+}
+
+function isIpRulesEnabled() {
+  return userContext.databaseAccount?.properties?.ipRules?.length > 0;
+}
+
+function isPrivateEndpointConnectionsEnabled() {
+  return userContext.databaseAccount?.properties?.privateEndpointConnections?.length > 0;
+}
+
+export function isPublicInternetAccessAllowed(): boolean {
+  return !isVirtualNetworkFilterEnabled() && !isIpRulesEnabled() && !isPrivateEndpointConnectionsEnabled();
+}
diff --git a/src/Contracts/DataModels.ts b/src/Contracts/DataModels.ts
index 1fb7d8674..e71c63cd9 100644
--- a/src/Contracts/DataModels.ts
+++ b/src/Contracts/DataModels.ts
@@ -22,6 +22,7 @@ export interface DatabaseAccountExtendedProperties {
   enableAnalyticalStorage?: boolean;
   isVirtualNetworkFilterEnabled?: boolean;
   ipRules?: IpRule[];
+  privateEndpointConnections?: unknown[];
 }
 
 export interface DatabaseAccountResponseLocation {
diff --git a/src/Explorer/Explorer.tsx b/src/Explorer/Explorer.tsx
index d40f4bcbd..fae2047ea 100644
--- a/src/Explorer/Explorer.tsx
+++ b/src/Explorer/Explorer.tsx
@@ -9,6 +9,7 @@ import * as Constants from "../Common/Constants";
 import { ExplorerMetrics } from "../Common/Constants";
 import { readCollection } from "../Common/dataAccess/readCollection";
 import { readDatabases } from "../Common/dataAccess/readDatabases";
+import { isPublicInternetAccessAllowed } from "../Common/DatabaseAccountUtility";
 import { getErrorMessage, getErrorStack, handleError } from "../Common/ErrorHandlingUtils";
 import * as Logger from "../Common/Logger";
 import { QueriesClient } from "../Common/QueriesClient";
@@ -188,11 +189,8 @@ export default class Explorer {
             ((await this._containsDefaultNotebookWorkspace(userContext.databaseAccount)) ||
               userContext.features.enableNotebooks)
         );
-        this.isShellEnabled(
-          this.isNotebookEnabled() &&
-            !userContext.databaseAccount.properties.isVirtualNetworkFilterEnabled &&
-            userContext.databaseAccount.properties.ipRules.length === 0
-        );
+
+        this.isShellEnabled(this.isNotebookEnabled() && isPublicInternetAccessAllowed());
 
         TelemetryProcessor.trace(Action.NotebookEnabled, ActionModifiers.Mark, {
           isNotebookEnabled: this.isNotebookEnabled(),
diff --git a/src/Explorer/Tree/ResourceTreeAdapter.tsx b/src/Explorer/Tree/ResourceTreeAdapter.tsx
index 3d7281acc..f06187564 100644
--- a/src/Explorer/Tree/ResourceTreeAdapter.tsx
+++ b/src/Explorer/Tree/ResourceTreeAdapter.tsx
@@ -14,6 +14,7 @@ import CollectionIcon from "../../../images/tree-collection.svg";
 import { ReactAdapter } from "../../Bindings/ReactBindingHandler";
 import { ArrayHashMap } from "../../Common/ArrayHashMap";
 import { Areas } from "../../Common/Constants";
+import { isPublicInternetAccessAllowed } from "../../Common/DatabaseAccountUtility";
 import * as DataModels from "../../Contracts/DataModels";
 import * as ViewModels from "../../Contracts/ViewModels";
 import { IPinnedRepo } from "../../Juno/JunoClient";
@@ -273,11 +274,7 @@ export class ResourceTreeAdapter implements ReactAdapter {
       contextMenu: ResourceTreeContextMenuButtonFactory.createCollectionContextMenuButton(this.container, collection),
     });
 
-    if (
-      userContext.apiType === "Mongo" &&
-      this.container.isNotebookEnabled() &&
-      userContext.features.enableSchemaAnalyzer
-    ) {
+    if (this.container.isNotebookEnabled() && userContext.apiType === "Mongo" && isPublicInternetAccessAllowed()) {
       children.push({
         label: "Schema (Preview)",
         onClick: collection.onSchemaAnalyzerClick.bind(collection),