Fix #6845 Notebook: Pyspark3 links incorrectly mapped when AD enabled (#6885)

src/sql/workbench/parts/notebook/common/models/cell.ts

@@ -24,6 +24,7 @@ import { generateUuid } from 'vs/base/common/uuid';
 import { IModelContentChangedEvent } from 'vs/editor/common/model/textModelEvents';
 let modelId = 0;
 
+
 export class CellModel implements ICellModel {
 	private _cellType: nb.CellType;
 	private _source: string | string[];
@@ -471,17 +472,20 @@ export class CellModel implements ICellModel {
 				if (result && result.data && result.data['text/html']) {
 					let model = (this as CellModel).options.notebook as NotebookModel;
 					if (model.activeConnection) {
-						let endpoint = this.getKnoxEndpoint(model.activeConnection);
+						let gatewayEndpointInfo = this.getGatewayEndpoint(model.activeConnection);
-						let host = endpoint && endpoint.ipAddress ? endpoint.ipAddress : model.activeConnection.serverName;
+						if (gatewayEndpointInfo) {
-						let port = endpoint && endpoint.port ? ':' + endpoint.port.toString() : defaultPort;
+							let hostAndIp = notebookUtils.getHostAndPortFromEndpoint(gatewayEndpointInfo.endpoint);
-						let html = result.data['text/html'];
+							let host = gatewayEndpointInfo && hostAndIp.host ? hostAndIp.host : model.activeConnection.serverName;
-						// CTP 3.1 and earlier Spark link
+							let port = gatewayEndpointInfo && hostAndIp.port ? ':' + hostAndIp.port : defaultPort;
-						html = this.rewriteUrlUsingRegex(/(https?:\/\/master.*\/proxy)(.*)/g, html, host, port, yarnUi);
+							let html = result.data['text/html'];
-						// CTP 3.2 and later spark link
+							// CTP 3.1 and earlier Spark link
-						html = this.rewriteUrlUsingRegex(/(https?:\/\/sparkhead.*\/proxy)(.*)/g, html, host, port, yarnUi);
+							html = this.rewriteUrlUsingRegex(/(https?:\/\/master.*\/proxy)(.*)/g, html, host, port, yarnUi);
-						// Driver link
+							// CTP 3.2 and later spark link
-						html = this.rewriteUrlUsingRegex(/(https?:\/\/storage.*\/containerlogs)(.*)/g, html, host, port, driverLog);
+							html = this.rewriteUrlUsingRegex(/(https?:\/\/sparkhead.*\/proxy)(.*)/g, html, host, port, yarnUi);
-						(<nb.IDisplayResult>output).data['text/html'] = html;
+							// Driver link
+							html = this.rewriteUrlUsingRegex(/(https?:\/\/storage.*\/containerlogs)(.*)/g, html, host, port, driverLog);
+							(<nb.IDisplayResult>output).data['text/html'] = html;
+						}
 					}
 				}
 			}
@@ -601,23 +605,21 @@ export class CellModel implements ICellModel {
 
 	// Get Knox endpoint from IConnectionProfile
 	// TODO: this will be refactored out into the notebooks extension as a contribution point
-	private getKnoxEndpoint(activeConnection: IConnectionProfile): notebookUtils.IEndpoint {
+	private getGatewayEndpoint(activeConnection: IConnectionProfile): notebookUtils.IEndpoint {
 		let endpoint;
 		if (this._connectionManagementService && activeConnection && activeConnection.providerName.toLowerCase() === notebookConstants.SQL_CONNECTION_PROVIDER.toLowerCase()) {
 			let serverInfo: ServerInfo = this._connectionManagementService.getServerInfo(activeConnection.id);
-			if (serverInfo && serverInfo.options && serverInfo.options['clusterEndpoints']) {
+			if (serverInfo) {
-				let endpoints: notebookUtils.IEndpoint[] = serverInfo.options['clusterEndpoints'];
+				let endpoints: notebookUtils.IEndpoint[] = notebookUtils.getClusterEndpoints(serverInfo);
 				if (endpoints && endpoints.length > 0) {
-					endpoint = endpoints.find(ep => {
+					endpoint = endpoints.find(ep => ep.serviceName.toLowerCase() === notebookUtils.hadoopEndpointNameGateway);
-						let serviceName: string = ep.serviceName.toLowerCase();
-						return serviceName === 'knox' || serviceName === 'gateway';
-					});
 				}
 			}
 		}
 		return endpoint;
 	}
 
+
 	private getMultilineSource(source: string | string[]): string | string[] {
 		if (typeof source === 'string') {
 			let sourceMultiline = source.split('\n');

src/sql/workbench/parts/notebook/common/models/notebookUtils.ts

@@ -4,13 +4,16 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as path from 'vs/base/common/path';
-import { nb } from 'azdata';
+import { nb, ServerInfo } from 'azdata';
 import { DEFAULT_NOTEBOOK_PROVIDER, DEFAULT_NOTEBOOK_FILETYPE, INotebookService } from 'sql/workbench/services/notebook/common/notebookService';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { ICellModel } from 'sql/workbench/parts/notebook/common/models/modelInterfaces';
+import { URI } from 'vs/base/common/uri';
 
 
+export const clusterEndpointsProperty = 'clusterEndpoints';
+export const hadoopEndpointNameGateway = 'gateway';
 /**
  * Test whether an output is from a stream.
  */
@@ -79,10 +82,12 @@ export interface IStandardKernelWithProvider {
 	readonly notebookProvider: string;
 }
 
+
 export interface IEndpoint {
 	serviceName: string;
-	ipAddress: string;
+	description: string;
-	port: number;
+	endpoint: string;
+	protocol: string;
 }
 
 export function tryMatchCellMagic(input: string): string {
@@ -125,3 +130,54 @@ export function convertVscodeResourceToFileInSubDirectories(htmlContent: string,
 export function useInProcMarkdown(configurationService: IConfigurationService): boolean {
 	return configurationService.getValue('notebook.useInProcMarkdown');
 }
+
+export function getClusterEndpoints(serverInfo: ServerInfo): IEndpoint[] | undefined {
+	let endpoints: RawEndpoint[] = serverInfo.options[clusterEndpointsProperty];
+	if (!endpoints || endpoints.length === 0) { return []; }
+
+	return endpoints.map(e => {
+		// If endpoint is missing, we're on CTP bits. All endpoints from the CTP serverInfo should be treated as HTTPS
+		let endpoint = e.endpoint ? e.endpoint : `https://${e.ipAddress}:${e.port}`;
+		let updatedEndpoint: IEndpoint = {
+			serviceName: e.serviceName,
+			description: e.description,
+			endpoint: endpoint,
+			protocol: e.protocol
+		};
+		return updatedEndpoint;
+	});
+}
+
+export type HostAndIp = { host: string, port: string };
+
+export function getHostAndPortFromEndpoint(endpoint: string): HostAndIp {
+	let authority = URI.parse(endpoint).authority;
+	let hostAndPortRegex = /^(.*)([,:](\d+))/g;
+	let match = hostAndPortRegex.exec(authority);
+	if (match) {
+		return {
+			host: match[1],
+			port: match[3]
+		};
+	}
+	return {
+		host: authority,
+		port: undefined
+	};
+}
+
+interface RawEndpoint {
+	serviceName: string;
+	description?: string;
+	endpoint?: string;
+	protocol?: string;
+	ipAddress?: string;
+	port?: number;
+}
+
+export interface IEndpoint {
+	serviceName: string;
+	description: string;
+	endpoint: string;
+	protocol: string;
+}