From 1e989060f91b56fe052bc6f1e27c9fd18d441112 Mon Sep 17 00:00:00 2001 From: Chris LaFreniere <40371649+chlafreniere@users.noreply.github.com> Date: Fri, 8 Mar 2019 17:34:47 -0800 Subject: [PATCH] Rewrite Spark UI link when using unified connection (#4362) * Rewrite Spark UI link when using unified connection * Add more robust error checking --- extensions/mssql/src/sqlClusterLookUp.ts | 7 +--- extensions/mssql/src/utils.ts | 2 +- extensions/notebook/src/common/utils.ts | 6 ++++ .../src/jupyter/jupyterSessionManager.ts | 12 ++----- .../parts/notebook/cellViews/codeActions.ts | 6 ++-- src/sql/parts/notebook/models/cell.ts | 34 ++++++++++++++++--- .../parts/notebook/models/modelInterfaces.ts | 2 +- src/sql/parts/notebook/notebook.component.ts | 2 +- src/sql/parts/notebook/notebookUtils.ts | 6 ++++ 9 files changed, 52 insertions(+), 25 deletions(-) diff --git a/extensions/mssql/src/sqlClusterLookUp.ts b/extensions/mssql/src/sqlClusterLookUp.ts index 354d227150..4bd77551f7 100644 --- a/extensions/mssql/src/sqlClusterLookUp.ts +++ b/extensions/mssql/src/sqlClusterLookUp.ts @@ -11,6 +11,7 @@ import * as UUID from 'vscode-languageclient/lib/utils/uuid'; import { AppContext } from './appContext'; import { SqlClusterConnection } from './objectExplorerNodeProvider/connection'; import { ICommandObjectExplorerContext } from './objectExplorerNodeProvider/command'; +import { IEndpoint } from './utils'; import { MssqlObjectExplorerNodeProvider } from './objectExplorerNodeProvider/objectExplorerNodeProvider'; @@ -118,12 +119,6 @@ function connToConnectionParam(connection: azdata.connection.Connection): Connec return result; } -interface IEndpoint { - serviceName: string; - ipAddress: string; - port: number; -} - class ConnectionParam implements azdata.connection.Connection, azdata.IConnectionProfile, azdata.ConnectionInfo { public connectionName: string; diff --git a/extensions/mssql/src/utils.ts b/extensions/mssql/src/utils.ts index 5b6a9f9da6..a1b3c54f3f 100644 --- a/extensions/mssql/src/utils.ts +++ b/extensions/mssql/src/utils.ts @@ -218,7 +218,7 @@ export async function getClusterEndpoint(profileId: string, serviceName: string) return clusterEndpoint; } -interface IEndpoint { +export interface IEndpoint { serviceName: string; ipAddress: string; port: number; diff --git a/extensions/notebook/src/common/utils.ts b/extensions/notebook/src/common/utils.ts index 2768b4f631..1678ca5edd 100644 --- a/extensions/notebook/src/common/utils.ts +++ b/extensions/notebook/src/common/utils.ts @@ -88,6 +88,12 @@ export enum Platform { Others } +export interface IEndpoint { + serviceName: string; + ipAddress: string; + port: number; +} + export function getOSPlatform(): Platform { switch (process.platform) { case 'win32': diff --git a/extensions/notebook/src/jupyter/jupyterSessionManager.ts b/extensions/notebook/src/jupyter/jupyterSessionManager.ts index 420159bc35..aca6d63d25 100644 --- a/extensions/notebook/src/jupyter/jupyterSessionManager.ts +++ b/extensions/notebook/src/jupyter/jupyterSessionManager.ts @@ -242,7 +242,7 @@ export class JupyterSession implements nb.ISession { //Update server info with bigdata endpoint - Unified Connection if (connection.providerName === SQL_PROVIDER) { - let clusterEndpoint: IEndpoint = await this.getClusterEndpoint(connection.id, KNOX_ENDPOINT); + let clusterEndpoint: utils.IEndpoint = await this.getClusterEndpoint(connection.id, KNOX_ENDPOINT); 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.'))); @@ -303,12 +303,12 @@ export class JupyterSession implements nb.ISession { return port; } - private async getClusterEndpoint(profileId: string, serviceName: string): Promise { + private async getClusterEndpoint(profileId: string, serviceName: string): Promise { let serverInfo: ServerInfo = await connection.getServerInfo(profileId); if (!serverInfo || !serverInfo.options) { return undefined; } - let endpoints: IEndpoint[] = serverInfo.options['clusterEndpoints']; + let endpoints: utils.IEndpoint[] = serverInfo.options['clusterEndpoints']; if (!endpoints || endpoints.length === 0) { return undefined; } @@ -320,12 +320,6 @@ interface ICredentials { 'url': string; } -interface IEndpoint { - serviceName: string; - ipAddress: string; - port: number; -} - interface ISparkMagicConfig { kernel_python_credentials: ICredentials; kernel_scala_credentials: ICredentials; diff --git a/src/sql/parts/notebook/cellViews/codeActions.ts b/src/sql/parts/notebook/cellViews/codeActions.ts index b5ad68b4f7..b93f874c92 100644 --- a/src/sql/parts/notebook/cellViews/codeActions.ts +++ b/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 { 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([ { 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 { 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); diff --git a/src/sql/parts/notebook/models/cell.ts b/src/sql/parts/notebook/models/cell.ts index 17182eceff..a1f3563d34 100644 --- a/src/sql/parts/notebook/models/cell.ts +++ b/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 { + public async runCell(notificationService?: INotificationService, connectionManagementService?: IConnectionManagementService): Promise { 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; + } } diff --git a/src/sql/parts/notebook/models/modelInterfaces.ts b/src/sql/parts/notebook/models/modelInterfaces.ts index 0ad6252555..7779654429 100644 --- a/src/sql/parts/notebook/models/modelInterfaces.ts +++ b/src/sql/parts/notebook/models/modelInterfaces.ts @@ -444,7 +444,7 @@ export interface ICellModel { readonly onExecutionStateChange: Event; setFuture(future: FutureInternal): void; readonly executionState: CellExecutionState; - runCell(notificationService?: INotificationService): Promise; + runCell(notificationService?: INotificationService, connectionManagementService?: IConnectionManagementService): Promise; setOverrideLanguage(language: string); equals(cellModel: ICellModel): boolean; toJSON(): nb.ICellContents; diff --git a/src/sql/parts/notebook/notebook.component.ts b/src/sql/parts/notebook/notebook.component.ts index ae2b4e3fb3..cc845a3a1f 100644 --- a/src/sql/parts/notebook/notebook.component.ts +++ b/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))); } diff --git a/src/sql/parts/notebook/notebookUtils.ts b/src/sql/parts/notebook/notebookUtils.ts index 297a61c63e..017bb349b4 100644 --- a/src/sql/parts/notebook/notebookUtils.ts +++ b/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;