diff --git a/src/sql/workbench/api/common/extHostNotebook.ts b/src/sql/workbench/api/common/extHostNotebook.ts index a7e48dbac7..3a6b48299e 100644 --- a/src/sql/workbench/api/common/extHostNotebook.ts +++ b/src/sql/workbench/api/common/extHostNotebook.ts @@ -173,6 +173,10 @@ export class ExtHostNotebook implements ExtHostNotebookShape { } $requestExecute(kernelId: number, content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): Thenable { + // Revive request's URIs to restore functions + content.notebookUri = URI.revive(content.notebookUri); + content.cellUri = URI.revive(content.cellUri); + let kernel = this._getAdapter(kernelId); let future = kernel.requestExecute(content, disposeOnDone); let futureId = this._addNewAdapter(future); @@ -259,9 +263,17 @@ export class ExtHostNotebook implements ExtHostNotebookShape { return this.registerSerializationProvider(serializationProvider); } - createNotebookController(extension: IExtensionDescription, id: string, viewType: string, label: string, handler?: (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable, rendererScripts?: vscode.NotebookRendererScript[]): vscode.NotebookController { + createNotebookController( + extension: IExtensionDescription, + id: string, + viewType: string, + label: string, + getDocHandler: (notebookUri: URI) => azdata.nb.NotebookDocument, + execHandler?: (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable, + rendererScripts?: vscode.NotebookRendererScript[] + ): vscode.NotebookController { let languagesHandler = (languages: string[]) => this._proxy.$updateKernelLanguages(viewType, viewType, languages); - let controller = new ADSNotebookController(extension, id, viewType, label, this._extHostNotebookDocumentsAndEditors, languagesHandler, handler, extension.enableProposedApi ? rendererScripts : undefined); + let controller = new ADSNotebookController(extension, id, viewType, label, this._extHostNotebookDocumentsAndEditors, languagesHandler, getDocHandler, execHandler, extension.enableProposedApi ? rendererScripts : undefined); let newKernel: azdata.nb.IStandardKernel = { name: viewType, displayName: controller.label, diff --git a/src/sql/workbench/api/common/notebooks/adsNotebookController.ts b/src/sql/workbench/api/common/notebooks/adsNotebookController.ts index 196bc094d4..ccd5ef7331 100644 --- a/src/sql/workbench/api/common/notebooks/adsNotebookController.ts +++ b/src/sql/workbench/api/common/notebooks/adsNotebookController.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import type * as vscode from 'vscode'; +import type * as azdata from 'azdata'; import { IExtensionDescription } from 'vs/platform/extensions/common/extensions'; import { INotebookKernelDto2 } from 'vs/workbench/api/common/extHost.protocol'; import { Emitter, Event } from 'vs/base/common/event'; @@ -14,13 +15,14 @@ import { URI } from 'vs/base/common/uri'; import { NotebookCellExecutionTaskState } from 'vs/workbench/api/common/extHostNotebookKernels'; import { asArray } from 'vs/base/common/arrays'; import { convertToADSCellOutput } from 'sql/workbench/api/common/notebooks/notebookUtils'; -import { CancellationToken } from 'vs/base/common/cancellation'; +import { CancellationTokenSource } from 'vs/base/common/cancellation'; type SelectionChangedEvent = { selected: boolean, notebook: vscode.NotebookDocument; }; type MessageReceivedEvent = { editor: vscode.NotebookEditor, message: any; }; type ExecutionHandler = (cells: vscode.NotebookCell[], notebook: vscode.NotebookDocument, controller: vscode.NotebookController) => void | Thenable; type LanguagesHandler = (languages: string[]) => void; type InterruptHandler = (notebook: vscode.NotebookDocument) => void | Promise; +type GetDocHandler = (notebookUri: URI) => azdata.nb.NotebookDocument; /** * A VS Code Notebook Controller that is used as part of converting VS Code notebook extension APIs into ADS equivalents. @@ -35,6 +37,8 @@ export class ADSNotebookController implements vscode.NotebookController { private readonly _languagesAdded = new Deferred(); private readonly _executionHandlerAdded = new Deferred(); + private readonly _execMap: Map = new Map(); + constructor( private _extension: IExtensionDescription, private _id: string, @@ -42,7 +46,8 @@ export class ADSNotebookController implements vscode.NotebookController { private _label: string, private _extHostNotebookDocumentsAndEditors: ExtHostNotebookDocumentsAndEditors, private _languagesHandler: LanguagesHandler, - private _handler?: ExecutionHandler, + private _getDocHandler: GetDocHandler, + private _execHandler?: ExecutionHandler, preloads?: vscode.NotebookRendererScript[] ) { this._kernelData = { @@ -53,7 +58,7 @@ export class ADSNotebookController implements vscode.NotebookController { label: this._label || this._extension.identifier.value, preloads: preloads ? preloads.map(extHostTypeConverters.NotebookRendererScript.from) : [] }; - if (this._handler) { + if (this._execHandler) { this._executionHandlerAdded.resolve(); } } @@ -125,11 +130,11 @@ export class ADSNotebookController implements vscode.NotebookController { } public get executeHandler(): ExecutionHandler { - return this._handler; + return this._execHandler; } public set executeHandler(value: ExecutionHandler) { - this._handler = value; + this._execHandler = value; this._executionHandlerAdded.resolve(); } @@ -142,8 +147,22 @@ export class ADSNotebookController implements vscode.NotebookController { this._kernelData.supportsInterrupt = Boolean(value); } + public getCellExecution(cellUri: URI): ADSNotebookCellExecution | undefined { + return this._execMap.get(cellUri.toString()); + } + + public removeCellExecution(cellUri: URI): void { + this._execMap.delete(cellUri.toString()); + } + + public getNotebookDocument(notebookUri: URI): azdata.nb.NotebookDocument { + return this._getDocHandler(notebookUri); + } + public createNotebookCellExecution(cell: vscode.NotebookCell): vscode.NotebookCellExecution { - return new ADSNotebookCellExecution(cell, this._extHostNotebookDocumentsAndEditors); + let exec = new ADSNotebookCellExecution(cell, this._extHostNotebookDocumentsAndEditors); + this._execMap.set(cell.document.uri.toString(), exec); + return exec; } public dispose(): void { @@ -166,6 +185,7 @@ export class ADSNotebookController implements vscode.NotebookController { class ADSNotebookCellExecution implements vscode.NotebookCellExecution { private _executionOrder: number; private _state = NotebookCellExecutionTaskState.Init; + private _cancellationSource = new CancellationTokenSource(); constructor(private readonly _cell: vscode.NotebookCell, private readonly _extHostNotebookDocumentsAndEditors: ExtHostNotebookDocumentsAndEditors) { this._executionOrder = this._cell.executionSummary?.executionOrder ?? -1; } @@ -174,8 +194,12 @@ class ADSNotebookCellExecution implements vscode.NotebookCellExecution { return this._cell; } + public get tokenSource(): vscode.CancellationTokenSource { + return this._cancellationSource; + } + public get token(): vscode.CancellationToken { - return CancellationToken.None; + return this._cancellationSource.token; } public get executionOrder(): number { diff --git a/src/sql/workbench/api/common/notebooks/vscodeExecuteProvider.ts b/src/sql/workbench/api/common/notebooks/vscodeExecuteProvider.ts index 9701ab4b9f..a103e3cbf9 100644 --- a/src/sql/workbench/api/common/notebooks/vscodeExecuteProvider.ts +++ b/src/sql/workbench/api/common/notebooks/vscodeExecuteProvider.ts @@ -9,6 +9,8 @@ import { ADSNotebookController } from 'sql/workbench/api/common/notebooks/adsNot import * as nls from 'vs/nls'; import { convertToVSCodeNotebookCell } from 'sql/workbench/api/common/notebooks/notebookUtils'; import { CellTypes } from 'sql/workbench/services/notebook/common/contracts'; +import { VSCodeNotebookDocument } from 'sql/workbench/api/common/notebooks/vscodeNotebookDocument'; +import { URI } from 'vs/base/common/uri'; class VSCodeFuture implements azdata.nb.IFuture { private _inProgress = true; @@ -71,6 +73,7 @@ class VSCodeKernel implements azdata.nb.IKernel { private readonly _name: string; private readonly _info: azdata.nb.IInfoReply; private readonly _kernelSpec: azdata.nb.IKernelSpec; + private _activeRequest: azdata.nb.IExecuteRequest; constructor(private readonly _controller: ADSNotebookController, private readonly _options: azdata.nb.ISessionOptions) { this._id = this._options.kernelId ?? (VSCodeKernel.kernelId++).toString(); @@ -136,11 +139,20 @@ class VSCodeKernel implements azdata.nb.IKernel { return Promise.resolve(this.spec); } + private cleanUpActiveExecution(cellUri: URI) { + this._activeRequest = undefined; + this._controller.removeCellExecution(cellUri); + } + requestExecute(content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): azdata.nb.IFuture { + if (this._activeRequest) { + throw new Error(nls.localize('notebookMultipleRequestsError', "Cannot execute code cell. Another cell is currently being executed.")); + } let executePromise: Promise; if (this._controller.executeHandler) { let cell = convertToVSCodeNotebookCell(CellTypes.Code, content.cellIndex, content.cellUri, content.notebookUri, content.language ?? this._kernelSpec.language, content.code); - executePromise = Promise.resolve(this._controller.executeHandler([cell], cell.notebook, this._controller)); + this._activeRequest = content; + executePromise = Promise.resolve(this._controller.executeHandler([cell], cell.notebook, this._controller)).then(() => this.cleanUpActiveExecution(content.cellUri)); } else { executePromise = Promise.resolve(); } @@ -154,7 +166,16 @@ class VSCodeKernel implements azdata.nb.IKernel { } public async interrupt(): Promise { - return; + if (this._activeRequest) { + if (this._controller.interruptHandler) { + let doc = this._controller.getNotebookDocument(this._activeRequest.notebookUri); + await this._controller.interruptHandler.call(this._controller, new VSCodeNotebookDocument(doc)); + } else { + let exec = this._controller.getCellExecution(this._activeRequest.cellUri); + exec?.tokenSource.cancel(); + } + this.cleanUpActiveExecution(this._activeRequest.cellUri); + } } public async restart(): Promise { diff --git a/src/vs/workbench/api/common/extHost.api.impl.ts b/src/vs/workbench/api/common/extHost.api.impl.ts index ea05e07e71..e0edb4b0c2 100644 --- a/src/vs/workbench/api/common/extHost.api.impl.ts +++ b/src/vs/workbench/api/common/extHost.api.impl.ts @@ -1158,7 +1158,8 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor, ex const notebooks: typeof vscode.notebooks = { createNotebookController(id: string, notebookType: string, label: string, handler?, rendererScripts?: vscode.NotebookRendererScript[]) { // {{SQL CARBON EDIT}} Use our own notebooks - return extHostNotebook.createNotebookController(extension, id, notebookType, label, handler, extension.enableProposedApi ? rendererScripts : undefined); + let getDocHandler = (notebookUri: URI) => extHostNotebookDocumentsAndEditors.getDocument(notebookUri.toString())?.document; + return extHostNotebook.createNotebookController(extension, id, notebookType, label, getDocHandler, handler, extension.enableProposedApi ? rendererScripts : undefined); }, registerNotebookCellStatusBarItemProvider: (notebookType: string, provider: vscode.NotebookCellStatusBarItemProvider) => { // {{SQL CARBON EDIT}} Disable VS Code notebooks