diff --git a/src/sql/azdata.d.ts b/src/sql/azdata.d.ts index 9a0cd54d03..999223c20b 100644 --- a/src/sql/azdata.d.ts +++ b/src/sql/azdata.d.ts @@ -4902,6 +4902,8 @@ declare module 'azdata' { * * @param disposeOnDone - Whether to dispose of the future when done. * + * @param cellId - Cell id (used by queryRunner) + * * @returns A kernel future. * * #### Notes @@ -4916,7 +4918,7 @@ declare module 'azdata' { * * **See also:** [[IExecuteReply]] */ - requestExecute(content: IExecuteRequest, disposeOnDone?: boolean): IFuture; + requestExecute(content: IExecuteRequest, disposeOnDone?: boolean, cellUri?: string): IFuture; /** * Send a `complete_request` message. diff --git a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts index 49949e72d6..bdef717af5 100644 --- a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts +++ b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts @@ -41,6 +41,7 @@ import { IQueryManagementService } from 'sql/workbench/services/query/common/que import { values } from 'vs/base/common/collections'; import { URI } from 'vs/base/common/uri'; import { assign } from 'vs/base/common/objects'; +import { escape } from 'sql/base/common/strings'; @Component({ selector: GridOutputComponent.SELECTOR, @@ -372,7 +373,7 @@ export class DataResourceDataProvider implements IGridDataProvider { } public async convertAllData(result: ResultSetSummary): Promise { - // Querying 50 rows at a time. Querying large amount of rows will be slow and + // Querying 100 rows at a time. Querying large amount of rows will be slow and // affect table rendering since each time the user scrolls, getRowData is called. let numRows = 100; for (let i = 0; i < result.rowCount; i += 100) { diff --git a/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts b/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts index 6c3a9332ee..8b71471f5d 100644 --- a/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts +++ b/src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts @@ -778,7 +778,7 @@ suite('Cell Model', function (): void { test('Execute returns error status', async function (): Promise { mockKernel.setup(k => k.requiresConnection).returns(() => false); - mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { + mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { let replyMsg: nb.IExecuteReplyMsg = { content: { execution_count: 1, @@ -796,7 +796,7 @@ suite('Cell Model', function (): void { test('Execute returns abort status', async function (): Promise { mockKernel.setup(k => k.requiresConnection).returns(() => false); - mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { + mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { let replyMsg: nb.IExecuteReplyMsg = { content: { execution_count: 1, @@ -815,7 +815,7 @@ suite('Cell Model', function (): void { test('Execute throws exception', async function (): Promise { let testMsg = 'Test message'; mockKernel.setup(k => k.requiresConnection).returns(() => false); - mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { + mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { throw new Error(testMsg); }); diff --git a/src/sql/workbench/services/notebook/browser/models/cell.ts b/src/sql/workbench/services/notebook/browser/models/cell.ts index 708bb8f63b..3c891e9834 100644 --- a/src/sql/workbench/services/notebook/browser/models/cell.ts +++ b/src/sql/workbench/services/notebook/browser/models/cell.ts @@ -408,7 +408,7 @@ export class CellModel extends Disposable implements ICellModel { const future = kernel.requestExecute({ code: content, stop_on_error: true - }, false); + }, false, this._cellUri.toString()); this.setFuture(future as FutureInternal); this.fireExecutionStateChanged(); // For now, await future completion. Later we should just track and handle cancellation based on model notifications diff --git a/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts b/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts index d9b94dd1b0..47496b8cb5 100644 --- a/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts +++ b/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts @@ -23,7 +23,6 @@ import { isUndefinedOrNull } from 'vs/base/common/types'; import { ILanguageMagic } from 'sql/workbench/services/notebook/browser/notebookService'; import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfigurationService'; import { URI } from 'vs/base/common/uri'; -import { getUriPrefix, uriPrefixes } from 'sql/platform/connection/common/utils'; import { firstIndex } from 'vs/base/common/arrays'; import { startsWith } from 'vs/base/common/strings'; import { onUnexpectedError } from 'vs/base/common/errors'; @@ -159,7 +158,7 @@ export class SqlSession implements nb.ISession { } class SqlKernel extends Disposable implements nb.IKernel { - private _queryRunner: QueryRunner; + private _queryRunners: Map = new Map(); private _currentConnection: IConnectionProfile; private _currentConnectionProfile: ConnectionProfile; static kernelId: number = 0; @@ -168,7 +167,6 @@ class SqlKernel extends Disposable implements nb.IKernel { private _future: SQLFuture; private _executionCount: number = 0; private _magicToExecutorMap = new Map(); - private _connectionPath: string; constructor(private _path: string, @IConnectionManagementService private _connectionManagementService: IConnectionManagementService, @@ -182,7 +180,6 @@ class SqlKernel extends Disposable implements nb.IKernel { ) { super(); this.initMagics(); - this.setConnectionPath(); } private initMagics(): void { @@ -192,29 +189,6 @@ class SqlKernel extends Disposable implements nb.IKernel { } } - private setConnectionPath(): void { - if (this._path) { - let prefix = getUriPrefix(this._path); - if (!prefix || prefix === uriPrefixes.connection) { - this._connectionPath = uriPrefixes.notebook.concat(this._path); - } else if (prefix !== uriPrefixes.notebook) { - try { - let uri = URI.parse(this._path); - if (uri && uri.scheme) { - this._connectionPath = uri.toString().replace(uri.scheme, uriPrefixes.notebook); - } - } catch { - // Ignore exceptions from URI parsing - } finally { - // If _connectionPath hasn't been set yet, set _connectionPath to _path as a last resort - if (!this._connectionPath) { - this._connectionPath = this._path; - } - } - } - } - } - public get id(): string { if (this._id === undefined) { this._id = (SqlKernel.kernelId++).toString(); @@ -265,30 +239,33 @@ class SqlKernel extends Disposable implements nb.IKernel { public set connection(conn: IConnectionProfile) { this._currentConnection = conn; this._currentConnectionProfile = new ConnectionProfile(this._capabilitiesService, this._currentConnection); - this._queryRunner = undefined; + this._queryRunners.clear(); } getSpec(): Thenable { return Promise.resolve(notebookConstants.sqlKernelSpec); } - requestExecute(content: nb.IExecuteRequest, disposeOnDone?: boolean): nb.IFuture { + requestExecute(content: nb.IExecuteRequest, disposeOnDone?: boolean, cellUri?: string): nb.IFuture { let canRun: boolean = true; let code = this.getCodeWithoutCellMagic(content); - if (this._queryRunner) { + let queryRunnerUri = 'queryRunner-' + cellUri; + let queryRunner: QueryRunner | undefined = this._queryRunners.get(queryRunnerUri); + if (queryRunner) { // Cancel any existing query - if (this._future && !this._queryRunner.hasCompleted) { - this._queryRunner.cancelQuery().then(ok => undefined, error => this._errorMessageService.showDialog(Severity.Error, sqlKernelError, error)); + if (this._future && !queryRunner.hasCompleted) { + queryRunner.cancelQuery().then(ok => undefined, error => this._errorMessageService.showDialog(Severity.Error, sqlKernelError, error)); // TODO when we can just show error as an output, should show an "execution canceled" error in output this._future.handleDone().catch(err => onUnexpectedError(err)); } - this._queryRunner.runQuery(code).catch(err => onUnexpectedError(err)); + queryRunner.runQuery(code).catch(err => onUnexpectedError(err)); } else if (this._currentConnection && this._currentConnectionProfile) { - this._queryRunner = this._instantiationService.createInstance(QueryRunner, this._connectionPath); - this.queryManagementService.registerRunner(this._queryRunner, this._connectionPath); - this._connectionManagementService.connect(this._currentConnectionProfile, this._connectionPath).then((result) => { - this.addQueryEventListeners(this._queryRunner); - this._queryRunner.runQuery(code).catch(err => onUnexpectedError(err)); + queryRunner = this._instantiationService.createInstance(QueryRunner, queryRunnerUri); + this._queryRunners.set(queryRunnerUri, queryRunner); + this.queryManagementService.registerRunner(queryRunner, queryRunnerUri); + this._connectionManagementService.connect(this._currentConnectionProfile, queryRunnerUri).then((result) => { + this.addQueryEventListeners(queryRunner); + queryRunner.runQuery(code).catch(err => onUnexpectedError(err)); }).catch(err => onUnexpectedError(err)); } else { canRun = false; @@ -298,7 +275,7 @@ class SqlKernel extends Disposable implements nb.IKernel { // TODO verify this is "canonical" behavior let count = canRun ? ++this._executionCount : undefined; - this._future = new SQLFuture(this._queryRunner, count, this._configurationService, this.logService); + this._future = new SQLFuture(queryRunner, count, this._configurationService, this.logService); if (!canRun) { // Complete early this._future.handleDone(new Error(localize('connectionRequired', "A connection must be chosen to run notebook cells"))).catch(err => onUnexpectedError(err)); @@ -334,8 +311,8 @@ class SqlKernel extends Disposable implements nb.IKernel { interrupt(): Thenable { // TODO: figure out what to do with the QueryCancelResult - return this._queryRunner.cancelQuery().then((cancelResult) => { - }); + let runners = [...this._queryRunners.values()]; + return Promise.all(runners.map(queryRunner => queryRunner.cancelQuery())).then(); } private addQueryEventListeners(queryRunner: QueryRunner): void { @@ -372,15 +349,16 @@ class SqlKernel extends Disposable implements nb.IKernel { } public async disconnect(): Promise { - if (this._connectionPath) { - if (this._connectionManagementService.isConnected(this._connectionPath)) { + this._queryRunners.forEach(async (queryRunner: QueryRunner, uri: string) => { + if (this._connectionManagementService.isConnected(uri)) { try { - await this._connectionManagementService.disconnect(this._connectionPath); + await this._connectionManagementService.disconnect(uri); } catch (err) { this.logService.error(err); } + } - } + }); return; } }