From b3be1d79cd845b8199fa8f0a8243302933d74277 Mon Sep 17 00:00:00 2001
From: Gene Lee <gelee@microsoft.com>
Date: Tue, 9 Apr 2019 15:19:02 -0700
Subject: [PATCH] Add support for new endpoint key string 'gateway'  (#4954)

---
 extensions/mssql/src/constants.ts                        | 3 ++-
 extensions/mssql/src/sqlClusterLookUp.ts                 | 6 +++++-
 extensions/notebook/src/jupyter/jupyterSessionManager.ts | 7 +++++--
 src/sql/workbench/parts/notebook/models/cell.ts          | 5 ++++-
 4 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/extensions/mssql/src/constants.ts b/extensions/mssql/src/constants.ts
index ddc2a6507c..305733b738 100644
--- a/extensions/mssql/src/constants.ts
+++ b/extensions/mssql/src/constants.ts
@@ -13,7 +13,8 @@ export const extensionConfigSectionName = 'mssql';
 
 // DATA PROTOCOL VALUES ///////////////////////////////////////////////////////////
 export const mssqlClusterProviderName = 'mssqlCluster';
-export const hadoopKnoxEndpointName = 'Knox';
+export const hadoopEndpointNameKnox = 'Knox';
+export const hadoopEndpointNameGateway = 'gateway';
 export const protocolVersion = '1.0';
 export const hostPropName = 'host';
 export const userPropName = 'user';
diff --git a/extensions/mssql/src/sqlClusterLookUp.ts b/extensions/mssql/src/sqlClusterLookUp.ts
index 4bd77551f7..f8131417a8 100644
--- a/extensions/mssql/src/sqlClusterLookUp.ts
+++ b/extensions/mssql/src/sqlClusterLookUp.ts
@@ -79,7 +79,11 @@ async function createSqlClusterConnInfo(sqlConnInfo: azdata.IConnectionProfile |
 	let endpoints: IEndpoint[] = serverInfo.options[constants.clusterEndpointsProperty];
 	if (!endpoints || endpoints.length === 0) { return undefined; }
 
-	let index = endpoints.findIndex(ep => ep.serviceName === constants.hadoopKnoxEndpointName);
+	let index = endpoints.findIndex(ep => {
+		let serviceName: string = ep.serviceName.toLowerCase();
+		return serviceName === constants.hadoopEndpointNameKnox.toLowerCase() ||
+			serviceName === constants.hadoopEndpointNameGateway.toLowerCase();
+	});
 	if (index < 0) { return undefined; }
 
 	let credentials = await azdata.connection.getCredentials(connectionId);
diff --git a/extensions/notebook/src/jupyter/jupyterSessionManager.ts b/extensions/notebook/src/jupyter/jupyterSessionManager.ts
index aca6d63d25..350f4e708a 100644
--- a/extensions/notebook/src/jupyter/jupyterSessionManager.ts
+++ b/extensions/notebook/src/jupyter/jupyterSessionManager.ts
@@ -57,7 +57,8 @@ const configBase = {
 
 const KNOX_ENDPOINT_SERVER = 'host';
 const KNOX_ENDPOINT_PORT = 'knoxport';
-const KNOX_ENDPOINT = 'knox';
+const KNOX_ENDPOINT_KNOX = 'knox';
+const KNOX_ENDPOINT_GATEWAY = 'gateway';
 const SQL_PROVIDER = 'MSSQL';
 const USER = 'user';
 const DEFAULT_CLUSTER_USER_NAME = 'root';
@@ -242,7 +243,9 @@ export class JupyterSession implements nb.ISession {
 
 			//Update server info with bigdata endpoint - Unified Connection
 			if (connection.providerName === SQL_PROVIDER) {
-				let clusterEndpoint: utils.IEndpoint = await this.getClusterEndpoint(connection.id, KNOX_ENDPOINT);
+				let clusterEndpoint: utils.IEndpoint =
+					await this.getClusterEndpoint(connection.id, KNOX_ENDPOINT_KNOX) ||
+					await this.getClusterEndpoint(connection.id, KNOX_ENDPOINT_GATEWAY);
 				if (!clusterEndpoint) {
 					let kernelDisplayName: string = await this.getKernelDisplayName();
 					return Promise.reject(new Error(localize('connectionNotValid', 'Spark kernels require a connection to a SQL Server big data cluster master instance.')));
diff --git a/src/sql/workbench/parts/notebook/models/cell.ts b/src/sql/workbench/parts/notebook/models/cell.ts
index c169d01143..cbb7ea265b 100644
--- a/src/sql/workbench/parts/notebook/models/cell.ts
+++ b/src/sql/workbench/parts/notebook/models/cell.ts
@@ -481,7 +481,10 @@ export class CellModel implements ICellModel {
 			if (serverInfo && serverInfo.options && serverInfo.options['clusterEndpoints']) {
 				let endpoints: notebookUtils.IEndpoint[] = serverInfo.options['clusterEndpoints'];
 				if (endpoints && endpoints.length > 0) {
-					endpoint = endpoints.find(ep => ep.serviceName.toLowerCase() === 'knox');
+					endpoint = endpoints.find(ep => {
+						let serviceName: string = ep.serviceName.toLowerCase();
+						return serviceName === 'knox' || serviceName === 'gateway';
+					});
 				}
 			}
 		}