diff --git a/extensions/integration-tests/src/notebook.test.ts b/extensions/integration-tests/src/notebook.test.ts index 3c9a0a125d..2613dc8585 100644 --- a/extensions/integration-tests/src/notebook.test.ts +++ b/extensions/integration-tests/src/notebook.test.ts @@ -15,6 +15,7 @@ import { getBdcServer, getConfigValue, EnvironmentVariable_PYTHON_PATH } from '. import { connectToServer, sleep } from './utils'; import * as fs from 'fs'; import { stressify } from 'adstest'; +import { isNullOrUndefined } from 'util'; if (context.RunTest) { suite('Notebook integration test suite', function () { @@ -36,6 +37,10 @@ if (context.RunTest) { await (new NotebookTester()).sqlNbMultipleCellsTest(this.test.title); }); + test('Sql NB run cells above and below test', async function () { + await (new NotebookTester()).sqlNbRunCellsAboveBelowTest(this.test.title); + }); + test('Clear cell output - SQL notebook', async function () { await (new NotebookTester()).sqlNbClearOutputs(this.test.title); }); @@ -108,6 +113,7 @@ class NotebookTester { @stressify({ dop: NotebookTester.ParallelCount }) async pySpark3NbTest(title: string): Promise { let notebook = await this.openNotebook(pySparkNotebookContent, pySpark3KernelMetadata, title + this.invocationCount++); + await this.runCell(notebook); let cellOutputs = notebook.document.cells[0].contents.outputs; let sparkResult = (cellOutputs[3]).text; assert(sparkResult === '2', `Expected spark result: 2, Actual: ${sparkResult}`); @@ -116,12 +122,14 @@ class NotebookTester { @stressify({ dop: NotebookTester.ParallelCount }) async python3ClearAllOutputs(title: string): Promise { let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title + this.invocationCount++); + await this.runCell(notebook); await this.verifyClearAllOutputs(notebook); } @stressify({ dop: NotebookTester.ParallelCount }) async python3NbTest(title: string): Promise { let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title + this.invocationCount++); + await this.runCell(notebook); let cellOutputs = notebook.document.cells[0].contents.outputs; console.log('Got cell outputs ---'); if (cellOutputs) { @@ -134,17 +142,20 @@ class NotebookTester { @stressify({ dop: NotebookTester.ParallelCount }) async sqlNbClearAllOutputs(title: string): Promise { let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++); + await this.runCell(notebook); await this.verifyClearAllOutputs(notebook); } async sqlNbClearOutputs(title: string): Promise { let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++); + await this.runCell(notebook); await this.verifyClearOutputs(notebook); } @stressify({ dop: NotebookTester.ParallelCount }) async sqlNbMultipleCellsTest(title: string): Promise { - let notebook = await this.openNotebook(sqlNotebookMultipleCellsContent, sqlKernelMetadata, title + this.invocationCount++, true); + let notebook = await this.openNotebook(sqlNotebookMultipleCellsContent, sqlKernelMetadata, title + this.invocationCount++); + await this.runCells(notebook); const expectedOutput0 = '(1 row affected)'; for (let i = 0; i < 3; i++) { let cellOutputs = notebook.document.cells[i].contents.outputs; @@ -162,9 +173,27 @@ class NotebookTester { } } + async sqlNbRunCellsAboveBelowTest(title: string): Promise { + let notebook = await this.openNotebook(sqlNotebookMultipleCellsContent, sqlKernelMetadata, title + this.invocationCount++); + // When running all cells above a cell, ensure that only cells preceding current cell have output + await this.runCells(notebook, true, undefined, notebook.document.cells[1]); + assert(notebook.document.cells[0].contents.outputs.length === 3, `Expected length: '3', Actual: '${notebook.document.cells[0].contents.outputs.length}'`); + assert(notebook.document.cells[1].contents.outputs.length === 0, `Expected length: '0', Actual: '${notebook.document.cells[1].contents.outputs.length}'`); + assert(notebook.document.cells[2].contents.outputs.length === 0, `Expected length: '0', Actual: '${notebook.document.cells[2].contents.outputs.length}'`); + + await notebook.clearAllOutputs(); + + // When running all cells below a cell, ensure that current cell and cells after have output + await this.runCells(notebook, undefined, true, notebook.document.cells[1]); + assert(notebook.document.cells[0].contents.outputs.length === 0, `Expected length: '0', Actual: '${notebook.document.cells[0].contents.outputs.length}'`); + assert(notebook.document.cells[1].contents.outputs.length === 3, `Expected length: '3', Actual: '${notebook.document.cells[1].contents.outputs.length}'`); + assert(notebook.document.cells[2].contents.outputs.length === 3, `Expected length: '3', Actual: '${notebook.document.cells[2].contents.outputs.length}'`); + } + @stressify({ dop: NotebookTester.ParallelCount }) async sqlNbTest(title: string): Promise { - let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++, false, true); + let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title + this.invocationCount++, true); + await this.runCell(notebook); const expectedOutput0 = '(1 row affected)'; let cellOutputs = notebook.document.cells[0].contents.outputs; console.log('Got cell outputs ---'); @@ -181,6 +210,7 @@ class NotebookTester { async sqlNbChangeKernelDifferentProviderTest(title: string): Promise { let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title); + await this.runCell(notebook); assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`); assert(notebook.document.kernelSpec.name === 'SQL', `Expected first kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`); @@ -196,6 +226,7 @@ class NotebookTester { async shouldNotBeDirtyAfterSavingNotebookTest(title: string): Promise { // Given a notebook that's been edited (in this case, open notebook runs the 1st cell and adds an output) let notebook = await this.openNotebook(sqlNotebookContent, sqlKernelMetadata, title); + await this.runCell(notebook); assert(notebook.document.providerId === 'sql', `Expected providerId to be sql, Actual: ${notebook.document.providerId}`); assert(notebook.document.kernelSpec.name === 'SQL', `Expected first kernel name: SQL, Actual: ${notebook.document.kernelSpec.name}`); assert(notebook.document.isDirty === true, 'Notebook should be dirty after edit'); @@ -230,6 +261,7 @@ class NotebookTester { async pythonChangeKernelDifferentProviderTest(title: string): Promise { let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title); + await this.runCell(notebook); assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`); assert(notebook.document.kernelSpec.name === 'python3', `Expected first kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`); @@ -244,6 +276,7 @@ class NotebookTester { async pythonChangeKernelSameProviderTest(title: string): Promise { let notebook = await this.openNotebook(pySparkNotebookContent, pythonKernelMetadata, title); + await this.runCell(notebook); assert(notebook.document.providerId === 'jupyter', `Expected providerId to be jupyter, Actual: ${notebook.document.providerId}`); assert(notebook.document.kernelSpec.name === 'python3', `Expected first kernel name: python3, Actual: ${notebook.document.kernelSpec.name}`); @@ -348,7 +381,7 @@ class NotebookTester { } } - async openNotebook(content: azdata.nb.INotebookContents, kernelMetadata: any, testName: string, runAllCells?: boolean, connectToDifferentServer?: boolean): Promise { + async openNotebook(content: azdata.nb.INotebookContents, kernelMetadata: any, testName: string, connectToDifferentServer?: boolean): Promise { let notebookConfig = vscode.workspace.getConfiguration('notebook'); notebookConfig.update('pythonPath', getConfigValue(EnvironmentVariable_PYTHON_PATH), 1); if (!connectToDifferentServer) { @@ -360,23 +393,30 @@ class NotebookTester { let uri = writeNotebookToFile(notebookJson, testName); console.log('Notebook uri ' + uri); let notebook = await azdata.nb.showNotebookDocument(uri); - console.log('Notebook is opened'); - - if (!runAllCells) { - assert(notebook.document.cells.length === 1, 'Notebook should have 1 cell'); - console.log('Before run notebook cell'); - let ran = await notebook.runCell(notebook.document.cells[0]); - console.log('After run notebook cell'); - assert(ran, 'Notebook runCell should succeed'); - } else { - console.log('Before run all notebook cells'); - let ran = await notebook.runAllCells(); - assert(ran, 'Notebook runCell should succeed'); - assert(notebook !== undefined && notebook !== null, 'Expected notebook object is defined'); - } - return notebook; } + + async runCells(notebook: azdata.nb.NotebookEditor, runCellsAbove?: boolean, runCellsBelow?: boolean, currentCell?: azdata.nb.NotebookCell) { + assert(notebook !== undefined && notebook !== null, 'Expected notebook object is defined'); + let ran; + if (runCellsAbove) { + ran = await notebook.runAllCells(undefined, currentCell); + } else if (runCellsBelow) { + ran = await notebook.runAllCells(currentCell, undefined); + } else { + ran = await notebook.runAllCells(); + } + assert(ran, 'Notebook runCell should succeed'); + } + + async runCell(notebook: azdata.nb.NotebookEditor, cell?: azdata.nb.NotebookCell) { + if (isNullOrUndefined(cell)) { + cell = notebook.document.cells[0]; + } + let ran = await notebook.runCell(cell); + assert(ran, 'Notebook runCell should succeed'); + } + async verifyClearAllOutputs(notebook: azdata.nb.NotebookEditor): Promise { let cellWithOutputs = notebook.document.cells.find(cell => cell.contents && cell.contents.outputs && cell.contents.outputs.length > 0); assert(cellWithOutputs !== undefined, 'Could not find notebook cells with outputs'); diff --git a/extensions/notebook/src/extension.ts b/extensions/notebook/src/extension.ts index 6169bd67e7..3d6987c595 100644 --- a/extensions/notebook/src/extension.ts +++ b/extensions/notebook/src/extension.ts @@ -176,11 +176,11 @@ async function clearActiveCellOutput(): Promise { } } -async function runAllCells(): Promise { +async function runAllCells(startCell?: azdata.nb.NotebookCell, endCell?: azdata.nb.NotebookCell): Promise { try { let notebook = azdata.nb.activeNotebookEditor; if (notebook) { - await notebook.runAllCells(); + await notebook.runAllCells(startCell, endCell); } else { throw new Error(noNotebookVisible); } diff --git a/src/sql/azdata.proposed.d.ts b/src/sql/azdata.proposed.d.ts index 57175424ce..857618b07b 100644 --- a/src/sql/azdata.proposed.d.ts +++ b/src/sql/azdata.proposed.d.ts @@ -4416,7 +4416,7 @@ declare module 'azdata' { /** * Kicks off execution of all code cells. Thenable will resolve only when full execution of all cells is completed. */ - runAllCells(): Thenable; + runAllCells(startCell?: NotebookCell, endCell?: NotebookCell): Thenable; /** * Clears the outputs of the active code cell in a notebook. diff --git a/src/sql/workbench/api/node/extHostNotebookEditor.ts b/src/sql/workbench/api/node/extHostNotebookEditor.ts index c8e45a5659..534f60b75b 100644 --- a/src/sql/workbench/api/node/extHostNotebookEditor.ts +++ b/src/sql/workbench/api/node/extHostNotebookEditor.ts @@ -155,8 +155,10 @@ export class ExtHostNotebookEditor implements azdata.nb.NotebookEditor, IDisposa return this._proxy.$runCell(this._id, uri); } - public runAllCells(): Thenable { - return this._proxy.$runAllCells(this._id); + public runAllCells(startCell?: azdata.nb.NotebookCell, endCell?: azdata.nb.NotebookCell): Thenable { + let startCellUri = startCell ? startCell.uri : undefined; + let endCellUri = endCell ? endCell.uri : undefined; + return this._proxy.$runAllCells(this._id, startCellUri, endCellUri); } public clearOutput(cell: azdata.nb.NotebookCell): Thenable { diff --git a/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts b/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts index 8259ae21b8..a6522fdec7 100644 --- a/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts +++ b/src/sql/workbench/api/node/mainThreadNotebookDocumentsAndEditors.ts @@ -132,11 +132,11 @@ class MainThreadNotebookEditor extends Disposable { return this.editor.runCell(cell); } - public runAllCells(): Promise { + public runAllCells(startCell?: ICellModel, endCell?: ICellModel): Promise { if (!this.editor) { return Promise.resolve(false); } - return this.editor.runAllCells(); + return this.editor.runAllCells(startCell, endCell); } public clearOutput(cell: ICellModel): Promise { @@ -383,12 +383,22 @@ export class MainThreadNotebookDocumentsAndEditors extends Disposable implements return editor.runCell(cell); } - $runAllCells(id: string): Promise { + $runAllCells(id: string, startCellUri?: UriComponents, endCellUri?: UriComponents): Promise { let editor = this.getEditor(id); if (!editor) { return Promise.reject(disposed(`TextEditor(${id})`)); } - return editor.runAllCells(); + let startCell: ICellModel; + let endCell: ICellModel; + if (startCellUri) { + let uriString = URI.revive(startCellUri).toString(); + startCell = editor.cells.find(c => c.cellUri.toString() === uriString); + } + if (endCellUri) { + let uriString = URI.revive(endCellUri).toString(); + endCell = editor.cells.find(c => c.cellUri.toString() === uriString); + } + return editor.runAllCells(startCell, endCell); } $clearOutput(id: string, cellUri: UriComponents): Promise { diff --git a/src/sql/workbench/api/node/sqlExtHost.protocol.ts b/src/sql/workbench/api/node/sqlExtHost.protocol.ts index bb9faf2748..4a10753bc4 100644 --- a/src/sql/workbench/api/node/sqlExtHost.protocol.ts +++ b/src/sql/workbench/api/node/sqlExtHost.protocol.ts @@ -919,7 +919,7 @@ export interface MainThreadNotebookDocumentsAndEditorsShape extends IDisposable $tryShowNotebookDocument(resource: UriComponents, options: INotebookShowOptions): Promise; $tryApplyEdits(id: string, modelVersionId: number, edits: ISingleNotebookEditOperation[], opts: IUndoStopOptions): Promise; $runCell(id: string, cellUri: UriComponents): Promise; - $runAllCells(id: string): Promise; + $runAllCells(id: string, startCellUri?: UriComponents, endCellUri?: UriComponents): Promise; $clearOutput(id: string, cellUri: UriComponents): Promise; $clearAllOutputs(id: string): Promise; $changeKernel(id: string, kernel: azdata.nb.IKernelInfo): Promise; diff --git a/src/sql/workbench/parts/notebook/cellToggleMoreActions.ts b/src/sql/workbench/parts/notebook/cellToggleMoreActions.ts index 4039c1a8d2..39ad826a8e 100644 --- a/src/sql/workbench/parts/notebook/cellToggleMoreActions.ts +++ b/src/sql/workbench/parts/notebook/cellToggleMoreActions.ts @@ -18,6 +18,7 @@ import { NotebookModel } from 'sql/workbench/parts/notebook/models/notebookModel import { ToggleMoreWidgetAction } from 'sql/workbench/parts/dashboard/common/actions'; import { CellTypes, CellType } from 'sql/workbench/parts/notebook/models/contracts'; import { CellModel } from 'sql/workbench/parts/notebook/models/cell'; +import { INotebookService } from 'sql/workbench/services/notebook/common/notebookService'; export const HIDDEN_CLASS = 'actionhidden'; @@ -33,6 +34,8 @@ export class CellToggleMoreActions { instantiationService.createInstance(AddCellFromContextAction, 'codeAfter', localize('codeAfter', 'Insert Code After'), CellTypes.Code, true), instantiationService.createInstance(AddCellFromContextAction, 'markdownBefore', localize('markdownBefore', 'Insert Text Before'), CellTypes.Markdown, false), instantiationService.createInstance(AddCellFromContextAction, 'markdownAfter', localize('markdownAfter', 'Insert Text After'), CellTypes.Markdown, true), + instantiationService.createInstance(RunCellsAction, 'runAllBefore', localize('runAllBefore', "Run Cells Before"), false), + instantiationService.createInstance(RunCellsAction, 'runAllAfter', localize('runAllAfter', "Run Cells After"), true), instantiationService.createInstance(ClearCellOutputAction, 'clear', localize('clear', 'Clear Output')) ); } @@ -141,3 +144,41 @@ export class ClearCellOutputAction extends CellActionBase { } } + +export class RunCellsAction extends CellActionBase { + constructor(id: string, + label: string, + private isAfter: boolean, + @INotificationService notificationService: INotificationService, + @INotebookService private notebookService: INotebookService, + ) { + super(id, label, undefined, notificationService); + } + + public canRun(context: CellContext): boolean { + return context.cell && context.cell.cellType === CellTypes.Code; + } + + async doRun(context: CellContext): Promise { + try { + let cell = context.cell || context.model.activeCell; + if (cell) { + let editor = this.notebookService.findNotebookEditor(cell.notebookModel.notebookUri); + if (editor) { + if (this.isAfter) { + await editor.runAllCells(cell, undefined); + } else { + await editor.runAllCells(undefined, cell); + } + } + } + } catch (error) { + let message = getErrorMessage(error); + this.notificationService.notify({ + severity: Severity.Error, + message: message + }); + } + return Promise.resolve(); + } +} \ No newline at end of file diff --git a/src/sql/workbench/parts/notebook/notebook.component.ts b/src/sql/workbench/parts/notebook/notebook.component.ts index 0037b029f4..92775a5432 100644 --- a/src/sql/workbench/parts/notebook/notebook.component.ts +++ b/src/sql/workbench/parts/notebook/notebook.component.ts @@ -50,6 +50,7 @@ import { ILogService } from 'vs/platform/log/common/log'; import { ITextFileService } from 'vs/workbench/services/textfile/common/textfiles'; import { LabeledMenuItemActionItem, fillInActions } from 'vs/platform/actions/browser/menuEntryActionViewItem'; import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry'; +import { isUndefinedOrNull } from 'vs/base/common/types'; export const NOTEBOOK_SELECTOR: string = 'notebook-component'; @@ -523,11 +524,20 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe } } - public async runAllCells(): Promise { + public async runAllCells(startCell?: ICellModel, endCell?: ICellModel): Promise { await this.modelReady; let codeCells = this._model.cells.filter(cell => cell.cellType === CellTypes.Code); if (codeCells && codeCells.length) { - for (let i = 0; i < codeCells.length; i++) { + // For the run all cells scenario where neither startId not endId are provided, set defaults + let startIndex = 0; + let endIndex = codeCells.length; + if (!isUndefinedOrNull(startCell)) { + startIndex = codeCells.findIndex(c => c.id === startCell.id); + } + if (!isUndefinedOrNull(endCell)) { + endIndex = codeCells.findIndex(c => c.id === endCell.id); + } + for (let i = startIndex; i < endIndex; i++) { let cellStatus = await this.runCell(codeCells[i]); if (!cellStatus) { return Promise.reject(new Error(localize('cellRunFailed', "Run Cells failed - See error in output of the currently selected cell for more information."))); diff --git a/src/sql/workbench/services/notebook/common/notebookService.ts b/src/sql/workbench/services/notebook/common/notebookService.ts index 1eb0521419..7b8663d499 100644 --- a/src/sql/workbench/services/notebook/common/notebookService.ts +++ b/src/sql/workbench/services/notebook/common/notebookService.ts @@ -128,7 +128,7 @@ export interface INotebookEditor { isVisible(): boolean; executeEdits(edits: ISingleNotebookEditOperation[]): boolean; runCell(cell: ICellModel): Promise; - runAllCells(): Promise; + runAllCells(startCell?: ICellModel, endCell?: ICellModel): Promise; clearOutput(cell: ICellModel): Promise; clearAllOutputs(): Promise; }