Rewrite Spark UI link when using unified connection (#4362)

* Rewrite Spark UI link when using unified connection

* Add more robust error checking

src/sql/parts/notebook/cellViews/codeActions.ts

@@ -15,6 +15,7 @@ import { INotificationService, Severity } from 'vs/platform/notification/common/
 import { NotebookModel } from 'sql/parts/notebook/models/notebookModel';
 import { getErrorMessage } from 'sql/parts/notebook/notebookUtils';
 import { ICellModel, CellExecutionState } from 'sql/parts/notebook/models/modelInterfaces';
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { MultiStateAction, IMultiStateData, IActionStateData } from 'sql/parts/notebook/notebookActions';
 
 let notebookMoreActionMsg = localize('notebook.failed', "Please select active cell and try again");
@@ -68,7 +69,8 @@ export class RunCellAction extends MultiStateAction<CellExecutionState> {
 	public static LABEL = 'Run cell';
 	private _executionChangedDisposable: IDisposable;
 	private _context: CellContext;
-	constructor(context: CellContext, @INotificationService private notificationService: INotificationService) {
+	constructor(context: CellContext, @INotificationService private notificationService: INotificationService,
+		@IConnectionManagementService private connectionManagementService: IConnectionManagementService) {
 		super(RunCellAction.ID, new IMultiStateData<CellExecutionState>([
 			{ key: CellExecutionState.Hidden, value: { label: emptyExecutionCountLabel, className: '', tooltip: '', hideIcon: true }},
 			{ key: CellExecutionState.Stopped, value: { label: '', className: 'toolbarIconRun', tooltip: localize('runCell', 'Run cell') }},
@@ -89,7 +91,7 @@ export class RunCellAction extends MultiStateAction<CellExecutionState> {
 			return;
 		}
 		try {
-			await this._context.cell.runCell(this.notificationService);
+			await this._context.cell.runCell(this.notificationService, this.connectionManagementService);
 		} catch (error) {
 			let message = getErrorMessage(error);
 			this.notificationService.error(message);

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

@@ -6,18 +6,21 @@
 
 'use strict';
 
-import { nb } from 'azdata';
+import { nb, ServerInfo } from 'azdata';
 
 import { Event, Emitter } from 'vs/base/common/event';
 import { URI } from 'vs/base/common/uri';
 import { localize } from 'vs/nls';
 
-import { ICellModelOptions, IModelFactory, FutureInternal, CellExecutionState } from './modelInterfaces';
 import * as notebookUtils from '../notebookUtils';
 import { CellTypes, CellType, NotebookChangeType } from 'sql/parts/notebook/models/contracts';
-import { ICellModel } from 'sql/parts/notebook/models/modelInterfaces';
 import { NotebookModel } from 'sql/parts/notebook/models/notebookModel';
+import { ICellModel } from 'sql/parts/notebook/models/modelInterfaces';
+import { ICellModelOptions, IModelFactory, FutureInternal, CellExecutionState } from './modelInterfaces';
+import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
+import { MssqlProviderId } from 'sql/common/constants';
 import { Schemas } from 'vs/base/common/network';
 let modelId = 0;
 
@@ -38,6 +41,7 @@ export class CellModel implements ICellModel {
 	private _executionCount: number | undefined;
 	private _cellUri: URI;
 	public id: string;
+	private _connectionManagementService: IConnectionManagementService;
 
 	constructor(private factory: IModelFactory, cellData?: nb.ICellContents, private _options?: ICellModelOptions) {
 		this.id = `${modelId++}`;
@@ -181,8 +185,11 @@ export class CellModel implements ICellModel {
 		return CellExecutionState.Hidden;
 	}
 
-	public async runCell(notificationService?: INotificationService): Promise<boolean> {
+	public async runCell(notificationService?: INotificationService, connectionManagementService?: IConnectionManagementService): Promise<boolean> {
 		try {
+			if (connectionManagementService) {
+				this._connectionManagementService = connectionManagementService;
+			}
 			if (this.cellType !== CellTypes.Code) {
 				// TODO should change hidden state to false if we add support
 				// for this property
@@ -371,7 +378,8 @@ 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 host = model.activeConnection.serverName;
+						let endpoint = this.getKnoxEndpoint(model.activeConnection);
+						let host = endpoint && endpoint.ipAddress ? endpoint.ipAddress : model.activeConnection.serverName;
 						let html = result.data['text/html'];
 						html = html.replace(/(https?:\/\/mssql-master.*\/proxy)(.*)/g, function (a, b, c) {
 							let ret = '';
@@ -458,4 +466,20 @@ export class CellModel implements ICellModel {
 		// Use this to set the internal (immutable) and public (shared with extension) uri properties
 		this.cellUri = uri;
 	}
+
+	// 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 {
+		let endpoint;
+		if (this._connectionManagementService && activeConnection && activeConnection.providerName === MssqlProviderId) {
+			let serverInfo: ServerInfo = this._connectionManagementService.getServerInfo(activeConnection.id);
+			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');
+				}
+			}
+		}
+		return endpoint;
+	}
 }

src/sql/parts/notebook/models/modelInterfaces.ts

@@ -444,7 +444,7 @@ export interface ICellModel {
 	readonly onExecutionStateChange: Event<CellExecutionState>;
 	setFuture(future: FutureInternal): void;
 	readonly executionState: CellExecutionState;
-	runCell(notificationService?: INotificationService): Promise<boolean>;
+	runCell(notificationService?: INotificationService, connectionManagementService?: IConnectionManagementService): Promise<boolean>;
 	setOverrideLanguage(language: string);
 	equals(cellModel: ICellModel): boolean;
 	toJSON(): nb.ICellContents;

src/sql/parts/notebook/notebook.component.ts

@@ -552,7 +552,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		await this.modelReady;
 		let uriString = cell.cellUri.toString();
 		if (this._model.cells.findIndex(c => c.cellUri.toString() === uriString) > -1) {
-			return cell.runCell(this.notificationService);
+			return cell.runCell(this.notificationService, this.connectionManagementService);
 		} else {
 			return Promise.reject(new Error(localize('cellNotFound', 'cell with URI {0} was not found in this model', uriString)));
 		}

src/sql/parts/notebook/notebookUtils.ts

@@ -101,6 +101,12 @@ export interface IStandardKernelWithProvider {
 	readonly notebookProvider: string;
 }
 
+export interface IEndpoint {
+	serviceName: string;
+	ipAddress: string;
+	port: number;
+}
+
 export function tryMatchCellMagic(input: string): string {
 	if (!input) {
 		return input;