From f1e21ebe9dbbbf59c922cd8c8c77583241888556 Mon Sep 17 00:00:00 2001
From: Yurong He <linda0306_hyr@hotmail.com>
Date: Fri, 26 Apr 2019 20:25:55 -0700
Subject: [PATCH] Fix driverlog uri in output (#5212)

* Fixed #4813

* Changed to lowercase for comparision
---
 src/sql/workbench/parts/notebook/models/cell.ts | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/sql/workbench/parts/notebook/models/cell.ts b/src/sql/workbench/parts/notebook/models/cell.ts
index 7575e0e659..ff8a1c62fa 100644
--- a/src/sql/workbench/parts/notebook/models/cell.ts
+++ b/src/sql/workbench/parts/notebook/models/cell.ts
@@ -12,7 +12,7 @@ import { localize } from 'vs/nls';
 import * as notebookUtils from '../notebookUtils';
 import { CellTypes, CellType, NotebookChangeType } from 'sql/workbench/parts/notebook/models/contracts';
 import { NotebookModel } from 'sql/workbench/parts/notebook/models/notebookModel';
-import { ICellModel } from 'sql/workbench/parts/notebook/models/modelInterfaces';
+import { ICellModel, notebookConstants } from 'sql/workbench/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';
@@ -383,6 +383,9 @@ export class CellModel implements ICellModel {
 	}
 
 	private rewriteOutputUrls(output: nb.ICellOutput): nb.ICellOutput {
+		const driverLog = '/gateway/default/yarn/container';
+		const yarnUi = '/gateway/default/yarn/proxy';
+		const defaultPort = ':30433';
 		// Only rewrite if this is coming back during execution, not when loading from disk.
 		// A good approximation is that the model has a future (needed for execution)
 		if (this.future) {
@@ -393,9 +396,10 @@ export class CellModel implements ICellModel {
 					if (model.activeConnection) {
 						let endpoint = this.getKnoxEndpoint(model.activeConnection);
 						let host = endpoint && endpoint.ipAddress ? endpoint.ipAddress : model.activeConnection.serverName;
+						let port = endpoint && endpoint.port ? ':' + endpoint.port.toString() : defaultPort;
 						let html = result.data['text/html'];
-						html = this.rewriteUrlUsingRegex(/(https?:\/\/mssql-master.*\/proxy)(.*)/g, html, host);
-						html = this.rewriteUrlUsingRegex(/(https?:\/\/master.*master-svc.*\/proxy)(.*)/g, html, host);
+						html = this.rewriteUrlUsingRegex(/(https?:\/\/master.*\/proxy)(.*)/g, html, host, port, yarnUi);
+						html = this.rewriteUrlUsingRegex(/(https?:\/\/storage.*\/containerlogs)(.*)/g, html, host, port, driverLog);
 						(<nb.IDisplayResult>output).data['text/html'] = html;
 					}
 				}
@@ -405,11 +409,11 @@ export class CellModel implements ICellModel {
 		return output;
 	}
 
-	private rewriteUrlUsingRegex(regex: RegExp, html: string, host: string): string {
+	private rewriteUrlUsingRegex(regex: RegExp, html: string, host: string, port: string, target: string): string {
 		return html.replace(regex, function (a, b, c) {
 			let ret = '';
 			if (b !== '') {
-				ret = 'https://' + host + ':30443/gateway/default/yarn/proxy';
+				ret = 'https://' + host + port + target;
 			}
 			if (c !== '') {
 				ret = ret + c;
@@ -489,7 +493,7 @@ export class CellModel implements ICellModel {
 	// 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 === 'mssql') {
+		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']) {
 				let endpoints: notebookUtils.IEndpoint[] = serverInfo.options['clusterEndpoints'];