Notebooks: Persist Chart Data when Re-Executing Cell (#14512)

* empty chart

* Update chart data appropriately

* Adding tests

* wip, cleanup

* PR feedback

src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts

@@ -278,11 +278,13 @@ class DataResourceTable extends GridTableBase<any> {
 	}
 
 	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<any> {
 	public updateResultSet(resultSet: ResultSetSummary, rows: ICellValue[][]): void {
 		this._gridDataProvider.updateResultSet(resultSet, rows);
 		super.updateResult(resultSet);
+		this.updateChartData(resultSet?.rowCount, resultSet?.columnInfo?.length, this._gridDataProvider);
 	}
 }
 

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: <IChartOption>{
+							configEntry: '',
+							default: '',
+							type: ControlType.input,
+							label: '',
+							displayableOptions: [''],
+						}
+					}
+				}
+			]
+		};
+
+		let future = TypeMoq.Mock.ofType(EmptyFuture);
+		let onIopub: nb.MessageHandler<nb.IIOPubMessage>;
+		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[<any>'_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[<any>'_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[<any>'_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');
+	});
 });

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<boolean>();
 	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: (<QueryResultId>msg.metadata).batchId, id: (<QueryResultId>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'];
 }