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 83f740b960..4fc2838409 100644 --- a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts +++ b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts @@ -278,11 +278,13 @@ class DataResourceTable extends GridTableBase { } public updateChartData(rowCount: number, columnCount: number, gridDataProvider: IGridDataProvider): void { - gridDataProvider.getRowData(0, rowCount).then(result => { - let range = new Slick.Range(0, 0, rowCount - 1, columnCount - 1); - let columns = gridDataProvider.getColumnHeaders(range); - this._chart.setData(result.rows, columns); - }); + if (this.chartDisplayed) { + gridDataProvider.getRowData(0, rowCount).then(result => { + let range = new Slick.Range(0, 0, rowCount - 1, columnCount - 1); + let columns = gridDataProvider.getColumnHeaders(range); + this._chart.setData(result.rows, columns); + }); + } } private setChartOptions(options: IInsightOptions | undefined) { @@ -293,6 +295,7 @@ class DataResourceTable extends GridTableBase { public updateResultSet(resultSet: ResultSetSummary, rows: ICellValue[][]): void { this._gridDataProvider.updateResultSet(resultSet, rows); super.updateResult(resultSet); + this.updateChartData(resultSet?.rowCount, resultSet?.columnInfo?.length, this._gridDataProvider); } } 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 e19082e592..985264256b 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 @@ -25,6 +25,8 @@ import { URI } from 'vs/base/common/uri'; import { IModelContentChangedEvent } from 'vs/editor/common/model/textModelEvents'; import { INotificationService } from 'vs/platform/notification/common/notification'; import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService'; +import { ControlType, IChartOption } from 'sql/workbench/contrib/charts/browser/chartOptions'; +import { CellModel } from 'sql/workbench/services/notebook/browser/models/cell'; let instantiationService: IInstantiationService; @@ -1075,4 +1077,79 @@ suite('Cell Model', function (): void { let serializedCell = model.toJSON(); assert.deepEqual(serializedCell.attachments, undefined, 'JSON should not include attachments if attachments do not exist'); }); + + test('Should not have cache chart data after new cell created', async function () { + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Code, + source: '' + }; + let cellModel = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }) as CellModel; + assert.deepEqual(cellModel.previousChartState, [], 'New cell should have no previous chart state'); + }); + + test('Should not cache chart data after clear output', async function () { + let notebookModel = new NotebookModelStub({ + name: '', + version: '', + mimetype: '' + }); + let contents: nb.ICellContents = { + cell_type: CellTypes.Code, + source: '', + outputs: [ + { + output_type: 'execute_result', + metadata: { + azdata_chartOptions: { + configEntry: '', + default: '', + type: ControlType.input, + label: '', + displayableOptions: [''], + } + } + } + ] + }; + + let future = TypeMoq.Mock.ofType(EmptyFuture); + let onIopub: nb.MessageHandler; + future.setup(f => f.setIOPubHandler(TypeMoq.It.isAny())).callback((handler) => onIopub = handler); + + // When I create a cell + let cellModel = factory.createCell(contents, { notebook: notebookModel, isTrusted: false }) as CellModel; + assert.deepEqual(cellModel.previousChartState, [], 'New cell should have no previous chart state'); + + // When previous chart state exists + cellModel['_previousChartState'] = contents.outputs[0].metadata.azdata_chartOptions; + assert.deepEqual(cellModel.previousChartState, contents.outputs[0].metadata.azdata_chartOptions, 'Previous chart state should be returned as is'); + + // When cell outputs are cleared + cellModel.clearOutputs(); + assert.deepEqual(cellModel.previousChartState, [], 'Previous chart state should be erased after clearing outputs'); + + // Put previous chart state back + cellModel['_previousChartState'] = contents.outputs[0].metadata.azdata_chartOptions; + + // When source is changed + cellModel.source = 'newSource'; + + // When output is generated + cellModel.setFuture(future.object); + await onIopub.handle({ channel: 'iopub', content: { data: 'Hello' }, type: 'execute_reply', metadata: contents.outputs[0].metadata, header: { msg_type: 'execute_result' } }); + assert.deepEqual(cellModel.previousChartState, [], 'Previous chart state should not exist after cell source change'); + + // Put previous chart state back + cellModel['_previousChartState'] = contents.outputs[0].metadata.azdata_chartOptions; + + // When output is generated + cellModel.setFuture(future.object); + await onIopub.handle({ channel: 'iopub', content: { data: 'Hello' }, type: 'execute_reply', metadata: contents.outputs[0].metadata, header: { msg_type: 'execute_result' } }); + assert.deepEqual(cellModel.previousChartState, contents.outputs[0].metadata.azdata_chartOptions, 'Previous chart state should exist after output is generated'); + }); }); diff --git a/src/sql/workbench/services/notebook/browser/models/cell.ts b/src/sql/workbench/services/notebook/browser/models/cell.ts index ef0c0ce9ad..31a19971eb 100644 --- a/src/sql/workbench/services/notebook/browser/models/cell.ts +++ b/src/sql/workbench/services/notebook/browser/models/cell.ts @@ -30,6 +30,7 @@ import { IConfigurationService } from 'vs/platform/configuration/common/configur import { Disposable } from 'vs/base/common/lifecycle'; import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys'; import { IAdsTelemetryService } from 'sql/platform/telemetry/common/telemetry'; +import { IInsightOptions } from 'sql/workbench/common/editor/query/chartState'; let modelId = 0; const ads_execute_command = 'ads_execute_command'; @@ -80,7 +81,10 @@ export class CellModel extends Disposable implements ICellModel { private _isParameter: boolean; private _onParameterStateChanged = new Emitter(); private _isInjectedParameter: boolean; + private _previousChartState: IInsightOptions[] = []; + private _outputCounter = 0; // When re-executing the same cell, ensure that we apply chart options in the same order private _attachments: nb.ICellAttachments; + private _preventNextChartCache: boolean = false; constructor(cellData: nb.ICellContents, private _options: ICellModelOptions, @@ -276,6 +280,7 @@ export class CellModel extends Disposable implements ICellModel { this.cellSourceChanged = true; } this._modelContentChangedEvent = undefined; + this._preventNextChartCache = true; } public get modelContentChangedEvent(): IModelContentChangedEvent { @@ -491,6 +496,7 @@ export class CellModel extends Disposable implements ICellModel { if (!kernel) { return false; } + this._outputCounter = 0; this._telemetryService?.createActionEvent(TelemetryKeys.TelemetryView.Notebook, TelemetryKeys.NbTelemetryAction.RunCell) .withAdditionalProperties({ cell_language: kernel.name }) .send(); @@ -626,14 +632,22 @@ export class CellModel extends Disposable implements ICellModel { if (this._future) { this._future.dispose(); } - this.clearOutputs(); + this.clearOutputs(true); this._future = future; future.setReplyHandler({ handle: (msg) => this.handleReply(msg) }); future.setIOPubHandler({ handle: (msg) => this.handleIOPub(msg) }); future.setStdInHandler({ handle: (msg) => this.handleSdtIn(msg) }); } - - public clearOutputs(): void { + /** + * Clear outputs can be done as part of the "Clear Outputs" action on a cell or as part of running a cell + * @param runCellPending If a cell has been run + */ + public clearOutputs(runCellPending = false): void { + if (runCellPending) { + this.cacheChartStateIfExists(); + } else { + this.clearPreviousChartState(); + } this._outputs = []; this._outputsIdMap.clear(); this.fireOutputsChanged(); @@ -641,6 +655,10 @@ export class CellModel extends Disposable implements ICellModel { this.executionCount = undefined; } + public get previousChartState(): any[] { + return this._previousChartState; + } + private fireOutputsChanged(shouldScroll: boolean = false): void { let outputEvent: IOutputChangedEvent = { outputs: this.outputs, @@ -705,7 +723,14 @@ export class CellModel extends Disposable implements ICellModel { } } if (!added) { + if (this._previousChartState[this._outputCounter]) { + if (!output.metadata) { + output.metadata = {}; + } + output.metadata.azdata_chartOptions = this._previousChartState[this._outputCounter]; + } this._outputsIdMap.set(output, { batchId: (msg.metadata).batchId, id: (msg.metadata).id }); + this._outputCounter++; } break; case 'execute_result_update': @@ -1023,4 +1048,33 @@ export class CellModel extends Disposable implements ICellModel { })); } } + + /** + * Cache start state for any existing charts. + * This ensures that data can be passed to the grid output component when a cell is re-executed + */ + private cacheChartStateIfExists(): void { + this.clearPreviousChartState(); + // If a cell's source was changed, don't cache chart state + if (!this._preventNextChartCache) { + this._outputs?.forEach(o => { + if (dataResourceDataExists(o)) { + if (o.metadata?.azdata_chartOptions) { + this._previousChartState.push(o.metadata.azdata_chartOptions); + } else { + this._previousChartState.push(undefined); + } + } + }); + } + this._preventNextChartCache = false; + } + + private clearPreviousChartState(): void { + this._previousChartState = []; + } +} + +function dataResourceDataExists(metadata: nb.ICellOutput): boolean { + return metadata['data']?.['application/vnd.dataresource+json']; }