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');
+	});
 });