Notebooks: re-factor grid streaming (#12937)

* refactor grid streaming (convert to data first)

* change convertRowsToHtml method to return value

* remove griddataconversioncomplete checks

* send row data from STS to gridoutput component

* clean up code

* send data updates to cell model

* serialize cell output at the end of cell execution

* remove unused parameters

* update output contents instead of output reference

* remove unnecessary promise

* move azdata changes to proposed

* update comment

src/sql/azdata.d.ts

@@ -4701,9 +4701,11 @@ declare module 'azdata' {
 
 		export interface ICellOutput {
 			output_type: OutputTypeName;
-			metadata?: {
+			metadata?: ICellOutputMetadata;
-				azdata_chartOptions?: any;
+		}
-			};
+
+		export interface ICellOutputMetadata {
+			azdata_chartOptions?: any;
 		}
 
 		/**
@@ -4736,7 +4738,9 @@ declare module 'azdata' {
 			/**
 			 * Optional metadata, also a mime bundle
 			 */
-			metadata?: {};
+			metadata?: {
+				resultSet?: ResultSetSummary;
+			};
 		}
 		export interface IDisplayData extends IDisplayResult {
 			output_type: 'display_data';

src/sql/azdata.proposed.d.ts

@@ -66,8 +66,16 @@ declare module 'azdata' {
 
 		export interface IExecuteResult {
 			data: any;
-			batchId?: number;
+		}
-			id?: number;
+
+		export interface IExecuteResultUpdate {
+			output_type: string;
+			resultSet: ResultSetSummary;
+			data: any;
+		}
+
+		export interface ICellOutputMetadata {
+			resultSet?: ResultSetSummary;
 		}
 	}
 

src/sql/workbench/contrib/notebook/browser/cellViews/output.component.ts

@@ -40,9 +40,8 @@ export class OutputComponent extends CellView implements OnInit, AfterViewInit {
 	private _initialized: boolean = false;
 	private _activeCellId: string;
 	private _componentInstance: IMimeComponent;
-	private _batchId?: number;
+	private _batchId: number | undefined;
-	private _id?: number;
+	private _id: number | undefined;
-	private _queryRunnerUri?: string;
 	public errorText: string;
 
 	constructor(
@@ -104,18 +103,14 @@ export class OutputComponent extends CellView implements OnInit, AfterViewInit {
 		return this._componentInstance;
 	}
 
-	@Input() set batchId(value: number) {
+	@Input() set batchId(value: number | undefined) {
 		this._batchId = value;
 	}
 
-	@Input() set id(value: number) {
+	@Input() set id(value: number | undefined) {
 		this._id = value;
 	}
 
-	@Input() set queryRunnerUri(value: string) {
-		this._queryRunnerUri = value;
-	}
-
 	get trustedMode(): boolean {
 		return this._trusted;
 	}
@@ -188,11 +183,8 @@ export class OutputComponent extends CellView implements OnInit, AfterViewInit {
 			this._componentInstance.cellModel = this.cellModel;
 			this._componentInstance.cellOutput = this.cellOutput;
 			this._componentInstance.bundleOptions = options;
-			if (this._queryRunnerUri) {
+			this._componentInstance.batchId = this._batchId;
-				this._componentInstance.batchId = this._batchId;
+			this._componentInstance.id = this._id;
-				this._componentInstance.id = this._id;
-				this._componentInstance.queryRunnerUri = this._queryRunnerUri;
-			}
 			this._changeref.detectChanges();
 			let el = <HTMLElement>componentRef.location.nativeElement;
 

src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.html

@@ -6,7 +6,7 @@
 -->
 <div style="overflow: hidden; width: 100%; height: 100%; display: flex; flex-flow: column">
 	<div #outputarea link-handler [isTrusted]="isTrusted" [notebookUri]="notebookUri" class="notebook-output" style="flex: 0 0 auto;">
-		<output-component *ngFor="let output of cellModel.outputs" [cellOutput]="output" [trustedMode] = "cellModel.trustedMode" [cellModel]="cellModel" [activeCellId]="activeCellId" [batchId]="output.batchId" [id]="output.id" [queryRunnerUri]="output.queryRunnerUri">
+		<output-component *ngFor="let output of cellModel.outputs" [cellOutput]="output" [trustedMode] = "cellModel.trustedMode" [cellModel]="cellModel" [activeCellId]="activeCellId" [batchId]="output.metadata?.resultSet?.batchId" [id]="output.metadata?.resultSet?.id">
 		</output-component>
 	</div>
 </div>

src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.ts

@@ -38,7 +38,6 @@ export class OutputAreaComponent extends AngularDisposable implements OnInit {
 			this._register(this.cellModel.onOutputsChanged(e => {
 				if (!(this._changeRef['destroyed'])) {
 					this._changeRef.detectChanges();
-					this._changeRef.detach();
 					if (e && e.shouldScroll) {
 						this.setFocusAndScroll(this.outputArea.nativeElement);
 					}

src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts

@@ -10,7 +10,6 @@ import { IInstantiationService } from 'vs/platform/instantiation/common/instanti
 import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions';
 import { NotebookInput } from 'sql/workbench/contrib/notebook/browser/models/notebookInput';
 import { INotebookService } from 'sql/workbench/services/notebook/browser/notebookService';
-import { INotificationService } from 'vs/platform/notification/common/notification';
 
 export class FileNotebookInput extends NotebookInput {
 	public static ID: string = 'workbench.editorinputs.fileNotebookInput';
@@ -22,10 +21,9 @@ export class FileNotebookInput extends NotebookInput {
 		@ITextModelService textModelService: ITextModelService,
 		@IInstantiationService instantiationService: IInstantiationService,
 		@INotebookService notebookService: INotebookService,
-		@IExtensionService extensionService: IExtensionService,
+		@IExtensionService extensionService: IExtensionService
-		@INotificationService notificationService: INotificationService
 	) {
-		super(title, resource, textInput, textModelService, instantiationService, notebookService, extensionService, notificationService);
+		super(title, resource, textInput, textModelService, instantiationService, notebookService, extensionService);
 	}
 
 	public get textInput(): FileEditorInput {

src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts

@@ -33,9 +33,6 @@ import { FileEditorInput } from 'vs/workbench/contrib/files/common/editors/fileE
 import { BinaryEditorModel } from 'vs/workbench/common/editor/binaryEditorModel';
 import { NotebookFindModel } from 'sql/workbench/contrib/notebook/browser/find/notebookFindModel';
 import { onUnexpectedError } from 'vs/base/common/errors';
-import { INotification, INotificationService } from 'vs/platform/notification/common/notification';
-import Severity from 'vs/base/common/severity';
-import * as nls from 'vs/nls';
 import { NotebookModel } from 'sql/workbench/services/notebook/browser/models/notebookModel';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
@@ -231,8 +228,7 @@ export abstract class NotebookInput extends EditorInput {
 		@ITextModelService private textModelService: ITextModelService,
 		@IInstantiationService private instantiationService: IInstantiationService,
 		@INotebookService private notebookService: INotebookService,
-		@IExtensionService private extensionService: IExtensionService,
+		@IExtensionService private extensionService: IExtensionService
-		@INotificationService private notificationService: INotificationService
 	) {
 		super();
 		this._standardKernels = [];
@@ -301,16 +297,6 @@ export abstract class NotebookInput extends EditorInput {
 	}
 
 	async save(groupId: number, options?: ITextFileSaveOptions): Promise<IEditorInput | undefined> {
-		const conversionNotification: INotification = {
-			severity: Severity.Info,
-			message: nls.localize('convertingData', "Waiting for table data conversion to complete..."),
-			progress: {
-				infinite: true // Keep showing conversion notification until notificationHandle is closed
-			}
-		};
-		const notificationHandle = this.notificationService.notify(conversionNotification);
-		await this._model.getNotebookModel().gridDataConversionComplete;
-		notificationHandle.close();
 		this.updateModel();
 		let input = await this.textInput.save(groupId, options);
 		await this.setTrustForNewEditor(input);

src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts

@@ -10,7 +10,6 @@ import { IExtensionService } from 'vs/workbench/services/extensions/common/exten
 import { NotebookInput } from 'sql/workbench/contrib/notebook/browser/models/notebookInput';
 import { INotebookService } from 'sql/workbench/services/notebook/browser/notebookService';
 import { UntitledTextEditorInput } from 'vs/workbench/services/untitled/common/untitledTextEditorInput';
-import { INotificationService } from 'vs/platform/notification/common/notification';
 
 export class UntitledNotebookInput extends NotebookInput {
 	public static ID: string = 'workbench.editorinputs.untitledNotebookInput';
@@ -22,10 +21,9 @@ export class UntitledNotebookInput extends NotebookInput {
 		@ITextModelService textModelService: ITextModelService,
 		@IInstantiationService instantiationService: IInstantiationService,
 		@INotebookService notebookService: INotebookService,
-		@IExtensionService extensionService: IExtensionService,
+		@IExtensionService extensionService: IExtensionService
-		@INotificationService notificationService: INotificationService
 	) {
-		super(title, resource, textInput, textModelService, instantiationService, notebookService, extensionService, notificationService);
+		super(title, resource, textInput, textModelService, instantiationService, notebookService, extensionService);
 	}
 
 	public get textInput(): UntitledTextEditorInput {

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

@@ -11,10 +11,10 @@ import { IContextMenuService } from 'vs/platform/contextview/browser/contextView
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
-import { IDataResource, MaxTableRowsConfigName, NotebookConfigSectionName, IDataResourceSchema, IDataResourceFields, MAX_ROWS } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
+import { IDataResource } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
 import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfigurationService';
-import QueryRunner, { getEolString, shouldIncludeHeaders, shouldRemoveNewLines } from 'sql/workbench/services/query/common/queryRunner';
+import { getEolString, shouldIncludeHeaders, shouldRemoveNewLines } from 'sql/workbench/services/query/common/queryRunner';
-import { ResultSetSummary, ResultSetSubset, ICellValue, BatchSummary } from 'sql/workbench/services/query/common/query';
+import { ResultSetSummary, ResultSetSubset, ICellValue } from 'sql/workbench/services/query/common/query';
 import { INotificationService } from 'vs/platform/notification/common/notification';
 import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService';
 import { attachTableStyler } from 'sql/platform/theme/common/styler';
@@ -37,11 +37,9 @@ import { ToggleableAction } from 'sql/workbench/contrib/notebook/browser/noteboo
 import { IInsightOptions } from 'sql/workbench/common/editor/query/chartState';
 import { NotebookChangeType } from 'sql/workbench/services/notebook/common/contracts';
 import { ActionsOrientation } from 'vs/base/browser/ui/actionbar/actionbar';
-import { IQueryManagementService } from 'sql/workbench/services/query/common/queryManagement';
 import { values } from 'vs/base/common/collections';
 import { URI } from 'vs/base/common/uri';
 import { assign } from 'vs/base/common/objects';
-import { escape } from 'sql/base/common/strings';
 
 @Component({
 	selector: GridOutputComponent.SELECTOR,
@@ -57,17 +55,13 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 	private _cellOutput: azdata.nb.ICellOutput;
 	private _bundleOptions: MimeModel.IOptions;
 	private _table: DataResourceTable;
-	private _batchId: number;
+	private _batchId: number | undefined;
-	private _id: number;
+	private _id: number | undefined;
-	private _queryRunnerUri: string;
+	private _layoutCalledOnce: boolean = false;
-	private _queryRunner: QueryRunner;
-	private _configuredMaxRows: number = MAX_ROWS;
 
 	constructor(
 		@Inject(IInstantiationService) private instantiationService: IInstantiationService,
-		@Inject(IThemeService) private readonly themeService: IThemeService,
+		@Inject(IThemeService) private readonly themeService: IThemeService
-		@Inject(IConfigurationService) private configurationService: IConfigurationService,
-		@Inject(IQueryManagementService) private queryManagementService: IQueryManagementService
 	) {
 		super();
 	}
@@ -100,28 +94,22 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 		this._cellOutput = value;
 	}
 
-	@Input() set batchId(value: number) {
+	@Input() set batchId(value: number | undefined) {
 		this._batchId = value;
 	}
 
-	@Input() set id(value: number) {
+	@Input() set id(value: number | undefined) {
 		this._id = value;
 	}
 
-	@Input() set queryRunnerUri(value: string) {
-		this._queryRunnerUri = value;
-	}
-
 	ngOnInit() {
-		let config = this.configurationService.getValue(NotebookConfigSectionName);
+		if (this.cellModel) {
-		if (config) {
+			this._register(this.cellModel.onTableUpdated(e => {
-			let maxRows = config[MaxTableRowsConfigName] ? config[MaxTableRowsConfigName] : undefined;
+				if (e.resultSet.batchId === this._batchId && e.resultSet.id === this._id) {
-			if (maxRows && maxRows > 0) {
+					this.updateResult(e.resultSet, e.rows);
-				this._configuredMaxRows = maxRows;
+				}
-			}
+			}));
 		}
-		// When a saved notebook is opened, there is no query runner
-		this._queryRunner = this.queryManagementService.getRunner(this._queryRunnerUri);
 		this.renderGrid();
 	}
 
@@ -132,51 +120,30 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 		if (!this._table) {
 			let source = <IDataResource><any>this._bundleOptions.data[this.mimeType];
 			let state = new GridTableState(0, 0);
-			this._table = this.instantiationService.createInstance(DataResourceTable, this._batchId, this._id, this._queryRunner, source, this.cellModel, this.cellOutput, state);
+			this._table = this.instantiationService.createInstance(DataResourceTable, source, this.cellModel, this.cellOutput, state);
 			let outputElement = <HTMLElement>this.output.nativeElement;
 			outputElement.appendChild(this._table.element);
 			this._register(attachTableStyler(this._table, this.themeService));
 			this._table.onDidInsert();
 			this.layout();
-			if (this._queryRunner) {
-				this._register(this._queryRunner.onResultSetUpdate(resultSet => { this.updateResultSet(resultSet); }));
-				this._register(this._queryRunner.onBatchEnd(batch => { this.convertData(batch); }));
-			}
 			this._initialized = true;
 		}
 	}
 
-	updateResultSet(resultSet: ResultSetSummary | ResultSetSummary[]): void {
-		let resultsToUpdate: ResultSetSummary[];
-		if (!Array.isArray(resultSet)) {
-			resultsToUpdate = [resultSet];
-		} else {
-			resultsToUpdate = resultSet?.splice(0);
-		}
-		for (let set of resultsToUpdate) {
-			if (this._batchId === set.batchId && this._id === set.id) {
-				set.rowCount = set.rowCount > this._configuredMaxRows ? this._configuredMaxRows : set.rowCount;
-				this._table.updateResult(set);
-				this.layout();
-			}
-		}
-	}
-
-	convertData(batch: BatchSummary): void {
-		for (let set of batch.resultSetSummaries) {
-			if (set.batchId === this._batchId && set.id === this._id) {
-				set.rowCount = set.rowCount > this._configuredMaxRows ? this._configuredMaxRows : set.rowCount;
-				this._cellModel.addGridDataConversionPromise(this._table.convertData(set));
-			}
-		}
-	}
-
 	layout(): void {
 		if (this._table) {
 			let maxSize = Math.min(this._table.maximumSize, 500);
 			this._table.layout(maxSize);
 		}
 	}
+
+	updateResult(resultSet: ResultSetSummary, rows: ICellValue[][]): void {
+		this._table.updateResultSet(resultSet, rows);
+		if (!this._layoutCalledOnce) {
+			this.layout();
+			this._layoutCalledOnce = true;
+		}
+	}
 }
 
 class DataResourceTable extends GridTableBase<any> {
@@ -184,13 +151,8 @@ class DataResourceTable extends GridTableBase<any> {
 	private _gridDataProvider: DataResourceDataProvider;
 	private _chart: ChartView;
 	private _chartContainer: HTMLElement;
-	private _batchId: number;
-	private _id: number;
-	private _queryRunner: QueryRunner;
 
-	constructor(batchId: number,
+	constructor(
-		id: number,
-		queryRunner: QueryRunner,
 		source: IDataResource,
 		private cellModel: ICellModel,
 		private cellOutput: azdata.nb.ICellOutput,
@@ -202,10 +164,7 @@ class DataResourceTable extends GridTableBase<any> {
 		@IConfigurationService configurationService: IConfigurationService
 	) {
 		super(state, createResultSet(source), { actionOrientation: ActionsOrientation.HORIZONTAL }, contextMenuService, instantiationService, editorService, untitledEditorService, configurationService);
-		this._batchId = batchId;
+		this._gridDataProvider = this.instantiationService.createInstance(DataResourceDataProvider, source, this.resultSet, this.cellModel);
-		this._id = id;
-		this._queryRunner = queryRunner;
-		this._gridDataProvider = this.instantiationService.createInstance(DataResourceDataProvider, this._batchId, this._id, this._queryRunner, source, this.resultSet, this.cellModel);
 		this._chart = this.instantiationService.createInstance(ChartView, false);
 
 		if (!this.cellOutput.metadata) {
@@ -292,28 +251,17 @@ class DataResourceTable extends GridTableBase<any> {
 		this.cellModel.sendChangeToNotebook(NotebookChangeType.CellMetadataUpdated);
 	}
 
-	public convertData(set: ResultSetSummary): Promise<void> {
+	public updateResultSet(resultSet: ResultSetSummary, rows: ICellValue[][]): void {
-		return this._gridDataProvider.convertAllData(set);
+		this._gridDataProvider.updateResultSet(resultSet, rows);
-	}
-
-	public updateResult(resultSet: ResultSetSummary): void {
 		super.updateResult(resultSet);
-		this._gridDataProvider.updateResultSet(resultSet);
 	}
 }
 
 export class DataResourceDataProvider implements IGridDataProvider {
 	private _rows: ICellValue[][];
 	private _documentUri: string;
-	private _queryRunner: QueryRunner;
-	private _batchId: number;
-	private _id: number;
 	private _resultSet: ResultSetSummary;
-	private _data: any;
 	constructor(
-		batchId: number,
-		id: number,
-		queryRunner: QueryRunner,
 		source: IDataResource,
 		resultSet: ResultSetSummary,
 		private cellModel: ICellModel,
@@ -325,42 +273,10 @@ export class DataResourceDataProvider implements IGridDataProvider {
 		@IInstantiationService private _instantiationService: IInstantiationService,
 	) {
 		this._documentUri = this.cellModel.notebookModel.notebookUri.toString();
-		this._queryRunner = queryRunner;
-		this._batchId = batchId;
-		this._id = id;
 		this._resultSet = resultSet;
-		this.initializeData();
 		this.transformSource(source);
 	}
 
-	private initializeData(): void {
-		// Set up data resource
-		let columnsResources: IDataResourceSchema[] = [];
-		this._resultSet.columnInfo.forEach(column => {
-			columnsResources.push({ name: escape(column.columnName) });
-		});
-		let columnsFields: IDataResourceFields = { fields: columnsResources };
-		let dataResource = {
-			schema: columnsFields,
-			data: []
-		};
-		// Set up html table string
-		let htmlTable: string[] = new Array(3);
-		htmlTable[0] = '<table>';
-		let columnHeaders = '<tr>';
-		for (let column of this._resultSet.columnInfo) {
-			columnHeaders += `<th>${escape(column.columnName)}</th>`;
-		}
-		columnHeaders += '</tr>';
-		htmlTable[1] = columnHeaders;
-		htmlTable[2] = '</table>';
-
-		this._data = {
-			'application/vnd.dataresource+json': dataResource,
-			'text/html': htmlTable
-		};
-	}
-
 	private transformSource(source: IDataResource): void {
 		this._rows = source.data.map(row => {
 			let rowData: azdata.DbCellValue[] = [];
@@ -377,46 +293,21 @@ export class DataResourceDataProvider implements IGridDataProvider {
 		});
 	}
 
-	public updateResultSet(resultSet: ResultSetSummary): void {
+	public updateResultSet(resultSet: ResultSetSummary, rows: ICellValue[][]): void {
 		this._resultSet = resultSet;
-	}
+		this._rows = this._rows.concat(rows);
-
-	public async convertAllData(result: ResultSetSummary): Promise<void> {
-		// Querying 100 rows at a time. Querying large amount of rows will be slow and
-		// affect table rendering since each time the user scrolls, getRowData is called.
-		let numRows = 100;
-		for (let i = 0; i < result.rowCount; i += 100) {
-			if (i + 100 > result.rowCount) {
-				numRows += result.rowCount - i;
-			}
-			let rows = await this._queryRunner.getQueryRows(i, numRows, this._batchId, this._id);
-			this.convertData(rows);
-		}
-		this.cellModel.sendChangeToNotebook(NotebookChangeType.CellOutputUpdated);
-	}
-
-	private convertData(rows: ResultSetSubset): void {
-		let dataResourceRows = this.convertRowsToDataResource(rows);
-		let htmlStringArr = this.convertRowsToHtml(rows);
-		this._data['application/vnd.dataresource+json'].data = this._data['application/vnd.dataresource+json'].data.concat(dataResourceRows);
-		this._data['text/html'].splice(this._data['text/html'].length - 1, 0, ...htmlStringArr);
-		this.cellModel.updateOutputData(this._batchId, this._id, this._data);
 	}
 
 	getRowData(rowStart: number, numberOfRows: number): Thenable<ResultSetSubset> {
-		if (this._queryRunner) {
+		let rowEnd = rowStart + numberOfRows;
-			return this._queryRunner.getQueryRows(rowStart, numberOfRows, this._batchId, this._id);
+		if (rowEnd > this._rows.length) {
-		} else {
+			rowEnd = this._rows.length;
-			let rowEnd = rowStart + numberOfRows;
-			if (rowEnd > this._rows.length) {
-				rowEnd = this._rows.length;
-			}
-			let resultSubset: ResultSetSubset = {
-				rowCount: rowEnd - rowStart,
-				rows: this._rows.slice(rowStart, rowEnd)
-			};
-			return Promise.resolve(resultSubset);
 		}
+		let resultSubset: ResultSetSubset = {
+			rowCount: rowEnd - rowStart,
+			rows: this._rows.slice(rowStart, rowEnd)
+		};
+		return Promise.resolve(resultSubset);
 	}
 
 	async copyResults(selection: Slick.Range[], includeHeaders?: boolean): Promise<void> {
@@ -454,13 +345,8 @@ export class DataResourceDataProvider implements IGridDataProvider {
 	}
 
 	serializeResults(format: SaveFormat, selection: Slick.Range[]): Thenable<void> {
-		if (this._queryRunner) {
+		let serializer = this._instantiationService.createInstance(ResultSerializer);
-			selection = selection ? selection : [new Slick.Range(0, 0, this._resultSet.rowCount - 1, this._resultSet.columnInfo.length - 1)];
+		return serializer.handleSerialization(this._documentUri, format, (filePath) => this.doSerialize(serializer, filePath, format, selection));
-			return this._queryRunner.serializeResults(this._batchId, this._id, format, selection);
-		} else {
-			let serializer = this._instantiationService.createInstance(ResultSerializer);
-			return serializer.handleSerialization(this._documentUri, format, (filePath) => this.doSerialize(serializer, filePath, format, selection));
-		}
 	}
 
 	private doSerialize(serializer: ResultSerializer, filePath: URI, format: SaveFormat, selection: Slick.Range[]): Promise<SaveResultsResponse | undefined> {
@@ -527,29 +413,6 @@ export class DataResourceDataProvider implements IGridDataProvider {
 	private isSelected(selection: Slick.Range): boolean {
 		return (selection && !((selection.fromCell === selection.toCell) && (selection.fromRow === selection.toRow)));
 	}
-
-	private convertRowsToDataResource(subset: ResultSetSubset): any[] {
-		return subset.rows.map(row => {
-			let rowObject: { [key: string]: any; } = {};
-			row.forEach((val, index) => {
-				rowObject[index] = val.displayValue;
-			});
-			return rowObject;
-		});
-	}
-
-	private convertRowsToHtml(subset: ResultSetSubset): string[] {
-		let htmlStringArr = [];
-		for (const row of subset.rows) {
-			let rowData = '<tr>';
-			for (let columnIndex = 0; columnIndex < row.length; columnIndex++) {
-				rowData += `<td>${escape(row[columnIndex].displayValue)}</td>`;
-			}
-			rowData += '</tr>';
-			htmlStringArr.push(rowData);
-		}
-		return htmlStringArr;
-	}
 }
 
 

src/sql/workbench/contrib/notebook/browser/outputs/mimeRegistry.ts

@@ -25,7 +25,6 @@ export interface IMimeComponent {
 	cellOutput?: nb.ICellOutput;
 	batchId?: number;
 	id?: number;
-	queryRunnerUri?: string;
 	layout(): void;
 }
 

src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts

@@ -13,7 +13,7 @@ import * as uuid from 'uuid';
 import * as sinon from 'sinon';
 import { DataResourceDataProvider } from '../../browser/outputs/gridOutput.component';
 import { IDataResource } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
-import { ResultSetSubset, ResultSetSummary } from 'sql/workbench/services/query/common/query';
+import { ResultSetSummary } from 'sql/workbench/services/query/common/query';
 import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
 import { TestFileDialogService, TestEditorService } from 'vs/workbench/test/browser/workbenchTestServices';
 import { TestContextService } from 'vs/workbench/test/common/workbenchTestServices';
@@ -23,7 +23,6 @@ import { InstantiationService } from 'vs/platform/instantiation/common/instantia
 import { URI } from 'vs/base/common/uri';
 import { CellModel } from 'sql/workbench/services/notebook/browser/models/cell';
 import { createandLoadNotebookModel } from 'sql/workbench/contrib/notebook/test/browser/cellToolbarActions.test';
-import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
 
 export class TestSerializationProvider implements azdata.SerializationProvider {
 	providerId: string;
@@ -98,9 +97,6 @@ suite('Data Resource Data Provider', function () {
 		let tempFolderPath = path.join(os.tmpdir(), `TestDataResourceDataProvider_${uuid.v4()}`);
 		await fs.mkdir(tempFolderPath);
 		let dataResourceDataProvider = new DataResourceDataProvider(
-			0, // batchId
-			0, // id
-			undefined, // QueryRunner
 			source,
 			resultSet,
 			cellModel.object,
@@ -131,38 +127,4 @@ suite('Data Resource Data Provider', function () {
 		const withHeadersResult = await fs.readFile(withHeadersFile.fsPath);
 		assert.equal(withHeadersResult.toString(), 'col1 col2 \n1 2 \n3 4 \n', 'result data should include headers');
 	});
-
-	test('convertAllData correctly converts row data to mimetype and html', async function (): Promise<void> {
-		let resultSetSubset: ResultSetSubset = {
-			rowCount: 2,
-			rows: [[{ displayValue: '1' }, { displayValue: '2' }], [{ displayValue: '3' }, { displayValue: '4' }]]
-		};
-		let queryRunner: TypeMoq.Mock<QueryRunner> = TypeMoq.Mock.ofType(QueryRunner);
-		queryRunner.setup(x => x.getQueryRows(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(resultSetSubset));
-		let dataResourceDataProvider = new DataResourceDataProvider(
-			0, // batchId
-			0, // id
-			queryRunner.object,
-			source,
-			resultSet,
-			cellModel.object,
-			notificationService,
-			undefined, // IClipboardService
-			undefined, // IConfigurationService
-			undefined, // ITextResourcePropertiesService
-			serializationService,
-			instantiationService.object
-		);
-		let spy = sinon.spy(cellModel.object, 'updateOutputData');
-		let expectedData = {
-			'application/vnd.dataresource+json': {
-				data: [{ 0: '1', 1: '2' }, { 0: '3', 1: '4' }],
-				schema: { fields: [{ name: 'col1' }, { name: 'col2' }] }
-			},
-			'text/html': ['<table>', '<tr><th>col1</th><th>col2</th></tr>', '<tr><td>1</td><td>2</td></tr>', '<tr><td>3</td><td>4</td></tr>', '</table>']
-		};
-		await dataResourceDataProvider.convertAllData(resultSet);
-		sinon.assert.calledOnce(spy);
-		sinon.assert.calledWithExactly(spy, 0, 0, expectedData);
-	});
 });

src/sql/workbench/contrib/notebook/test/browser/notebookEditor.test.ts

@@ -56,8 +56,6 @@ import { UntitledTextEditorInput } from 'vs/workbench/services/untitled/common/u
 import { IUntitledTextEditorService } from 'vs/workbench/services/untitled/common/untitledTextEditorService';
 import { workbenchInstantiationService } from 'vs/workbench/test/browser/workbenchTestServices';
 import { IProductService } from 'vs/platform/product/common/productService';
-import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
-import { INotificationService } from 'vs/platform/notification/common/notification';
 
 class NotebookModelStub extends stubs.NotebookModelStub {
 	public contentChangedEmitter = new Emitter<NotebookContentChange>();
@@ -97,11 +95,10 @@ suite.skip('Test class NotebookEditor:', () => {
 	let queryTextEditor: QueryTextEditor;
 	let untitledNotebookInput: UntitledNotebookInput;
 	let notebookEditorStub: NotebookEditorStub;
-	let notificationService: TypeMoq.Mock<INotificationService>;
 
 	setup(async () => {
 		// setup services
-		({ instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub, notificationService } = setupServices({ instantiationService, workbenchThemeService }));
+		({ instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub } = setupServices({ instantiationService, workbenchThemeService }));
 		// Create notebookEditor
 		notebookEditor = createNotebookEditor(instantiationService, workbenchThemeService, notebookService);
 	});
@@ -122,7 +119,7 @@ suite.skip('Test class NotebookEditor:', () => {
 			const untitledTextInput = instantiationService.createInstance(UntitledTextEditorInput, untitledTextEditorService.create({ associatedResource: untitledUri }));
 			const untitledNotebookInput = new UntitledNotebookInput(
 				testTitle, untitledUri, untitledTextInput,
-				undefined, instantiationService, notebookService, extensionService, notificationService.object
+				undefined, instantiationService, notebookService, extensionService
 			);
 			const testNotebookEditor = new NotebookEditorStub({ cellGuid: cellTextEditorGuid, editor: queryTextEditor, model: notebookModel, notebookParams: <INotebookParams>{ notebookUri: untitledNotebookInput.notebookUri } });
 			notebookService.addNotebookEditor(testNotebookEditor);
@@ -670,7 +667,6 @@ function setupServices(arg: { workbenchThemeService?: WorkbenchThemeService, ins
 	const uninstallEvent = new Emitter<IExtensionIdentifier>();
 	const didUninstallEvent = new Emitter<DidUninstallExtensionEvent>();
 
-	const notificationService = TypeMoq.Mock.ofType(TestNotificationService, TypeMoq.MockBehavior.Loose);
 	const instantiationService = arg.instantiationService ?? <TestInstantiationService>workbenchInstantiationService();
 	const workbenchThemeService = arg.workbenchThemeService ?? instantiationService.createInstance(WorkbenchThemeService);
 	instantiationService.stub(IWorkbenchThemeService, workbenchThemeService);
@@ -707,7 +703,7 @@ function setupServices(arg: { workbenchThemeService?: WorkbenchThemeService, ins
 	const untitledTextInput = instantiationService.createInstance(UntitledTextEditorInput, untitledTextEditorService.create({ associatedResource: untitledUri }));
 	const untitledNotebookInput = new UntitledNotebookInput(
 		testTitle, untitledUri, untitledTextInput,
-		undefined, instantiationService, notebookService, extensionService, notificationService.object
+		undefined, instantiationService, notebookService, extensionService
 	);
 
 	const cellTextEditorGuid = generateUuid();
@@ -722,7 +718,7 @@ function setupServices(arg: { workbenchThemeService?: WorkbenchThemeService, ins
 	);
 	const notebookEditorStub = new NotebookEditorStub({ cellGuid: cellTextEditorGuid, editor: queryTextEditor, model: new NotebookModelStub(), notebookParams: <INotebookParams>{ notebookUri: untitledNotebookInput.notebookUri } });
 	notebookService.addNotebookEditor(notebookEditorStub);
-	return { instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub, notificationService };
+	return { instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub };
 }
 
 function createNotebookEditor(instantiationService: TestInstantiationService, workbenchThemeService: WorkbenchThemeService, notebookService: NotebookService) {

src/sql/workbench/contrib/notebook/test/browser/notebookInput.test.ts

@@ -20,7 +20,6 @@ import { IExtensionService, NullExtensionService } from 'vs/workbench/services/e
 import { INotebookService, IProviderInfo } from 'sql/workbench/services/notebook/browser/notebookService';
 import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock';
 import { IUntitledTextEditorService } from 'vs/workbench/services/untitled/common/untitledTextEditorService';
-import { TestNotificationService } from 'vs/platform/notification/test/common/testNotificationService';
 
 suite('Notebook Input', function (): void {
 	const instantiationService = workbenchInstantiationService();
@@ -45,8 +44,6 @@ suite('Notebook Input', function (): void {
 
 	(instantiationService as TestInstantiationService).stub(INotebookService, mockNotebookService.object);
 
-	const mockNotificationService = TypeMoq.Mock.ofType(TestNotificationService);
-
 	let untitledTextInput: UntitledTextEditorInput;
 	let untitledNotebookInput: UntitledNotebookInput;
 
@@ -56,14 +53,14 @@ suite('Notebook Input', function (): void {
 		untitledTextInput = instantiationService.createInstance(UntitledTextEditorInput, service.create({ associatedResource: untitledUri }));
 		untitledNotebookInput = new UntitledNotebookInput(
 			testTitle, untitledUri, untitledTextInput,
-			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object, mockNotificationService.object);
+			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object);
 	});
 
 	test('File Notebook Input', async function (): Promise<void> {
 		let fileUri = URI.from({ scheme: Schemas.file, path: 'TestPath' });
 		let fileNotebookInput = new FileNotebookInput(
 			testTitle, fileUri, undefined,
-			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object, mockNotificationService.object);
+			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object);
 
 		let inputId = fileNotebookInput.getTypeId();
 		assert.strictEqual(inputId, FileNotebookInput.ID);

src/sql/workbench/contrib/notebook/test/stubs.ts

@@ -43,9 +43,6 @@ export class NotebookModelStub implements INotebookModel {
 	get sessionLoadFinished(): Promise<void> {
 		throw new Error('method not implemented.');
 	}
-	get gridDataConversionComplete(): Promise<any[]> {
-		throw new Error('method not implemented.');
-	}
 	get notebookManagers(): INotebookManager[] {
 		throw new Error('method not implemented.');
 	}

src/sql/workbench/services/notebook/browser/models/cell.ts

@@ -12,7 +12,7 @@ import { localize } from 'vs/nls';
 import * as notebookUtils from 'sql/workbench/services/notebook/browser/models/notebookUtils';
 import { CellTypes, CellType, NotebookChangeType } from 'sql/workbench/services/notebook/common/contracts';
 import { NotebookModel } from 'sql/workbench/services/notebook/browser/models/notebookModel';
-import { ICellModel, IOutputChangedEvent, CellExecutionState, ICellModelOptions } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
+import { ICellModel, IOutputChangedEvent, CellExecutionState, ICellModelOptions, ITableUpdatedEvent } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
@@ -28,6 +28,7 @@ import { ICommandService } from 'vs/platform/commands/common/commands';
 import { tryMatchCellMagic, extractCellMagicCommandPlusArgs } from 'sql/workbench/services/notebook/browser/utils';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { Disposable } from 'vs/base/common/lifecycle';
+import { ResultSetSummary } from 'sql/workbench/services/query/common/query';
 
 let modelId = 0;
 const ads_execute_command = 'ads_execute_command';
@@ -44,6 +45,7 @@ export class CellModel extends Disposable implements ICellModel {
 	private _renderedOutputTextContent: string[] = [];
 	private _isEditMode: boolean;
 	private _onOutputsChanged = new Emitter<IOutputChangedEvent>();
+	private _onTableUpdated = new Emitter<ITableUpdatedEvent>();
 	private _onCellModeChanged = new Emitter<boolean>();
 	private _onExecutionStateChanged = new Emitter<CellExecutionState>();
 	private _isTrusted: boolean;
@@ -66,7 +68,6 @@ export class CellModel extends Disposable implements ICellModel {
 	private _showPreview: boolean = true;
 	private _showMarkdown: boolean = false;
 	private _cellSourceChanged: boolean = false;
-	private _gridDataConversionComplete: Promise<void>[] = [];
 	private _defaultToWYSIWYG: boolean;
 	private _isParameter: boolean;
 	private _onParameterStateChanged = new Emitter<boolean>();
@@ -113,6 +114,10 @@ export class CellModel extends Disposable implements ICellModel {
 		return this._onOutputsChanged.event;
 	}
 
+	public get onTableUpdated(): Event<ITableUpdatedEvent> {
+		return this._onTableUpdated.event;
+	}
+
 	public get onCellModeChanged(): Event<boolean> {
 		return this._onCellModeChanged.event;
 	}
@@ -441,8 +446,6 @@ export class CellModel extends Disposable implements ICellModel {
 
 	public async runCell(notificationService?: INotificationService, connectionManagementService?: IConnectionManagementService): Promise<boolean> {
 		try {
-			// Clear grid data conversion promises from previous execution results
-			this._gridDataConversionComplete = [];
 			if (!this.active && this !== this.notebookModel.activeCell) {
 				this.notebookModel.updateActiveCell(this);
 				this.active = true;
@@ -536,6 +539,8 @@ export class CellModel extends Disposable implements ICellModel {
 		} finally {
 			this.disposeFuture();
 			this.fireExecutionStateChanged();
+			// Serialize cell output once the cell is done executing
+			this.sendChangeToNotebook(NotebookChangeType.CellOutputUpdated);
 			this.notifyExecutionComplete();
 		}
 
@@ -609,9 +614,7 @@ export class CellModel extends Disposable implements ICellModel {
 			shouldScroll: !!shouldScroll
 		};
 		this._onOutputsChanged.fire(outputEvent);
-		if (this.outputs.length !== 0) {
+		if (this.outputs.length === 0) {
-			this.sendChangeToNotebook(NotebookChangeType.CellOutputUpdated);
-		} else {
 			this.sendChangeToNotebook(NotebookChangeType.CellOutputCleared);
 		}
 	}
@@ -626,25 +629,6 @@ export class CellModel extends Disposable implements ICellModel {
 		return this._outputs;
 	}
 
-	public updateOutputData(batchId: number, id: number, data: any) {
-		for (let i = 0; i < this._outputs.length; i++) {
-			if (this._outputs[i].output_type === 'execute_result'
-				&& (<nb.IExecuteResult>this._outputs[i]).batchId === batchId
-				&& (<nb.IExecuteResult>this._outputs[i]).id === id) {
-				(<nb.IExecuteResult>this._outputs[i]).data = data;
-				break;
-			}
-		}
-	}
-
-	public get gridDataConversionComplete(): Promise<void> {
-		return Promise.all(this._gridDataConversionComplete).then();
-	}
-
-	public addGridDataConversionPromise(complete: Promise<void>): void {
-		this._gridDataConversionComplete.push(complete);
-	}
-
 	public get renderedOutputTextContent(): string[] {
 		return this._renderedOutputTextContent;
 	}
@@ -665,9 +649,49 @@ export class CellModel extends Disposable implements ICellModel {
 	private handleIOPub(msg: nb.IIOPubMessage): void {
 		let msgType = msg.header.msg_type;
 		let output: nb.ICellOutput;
+		let added = false;
 		switch (msgType) {
 			case 'execute_result':
+				output = msg.content as nb.ICellOutput;
+				output.output_type = msgType;
+				// Check if the table already exists
+				for (let i = 0; i < this._outputs.length; i++) {
+					if (this._outputs[i].output_type === 'execute_result') {
+						let resultSet: ResultSetSummary = this._outputs[i].metadata.resultSet;
+						let newResultSet: ResultSetSummary = output.metadata.resultSet;
+						if (resultSet.batchId === newResultSet.batchId && resultSet.id === newResultSet.id) {
+							// If it does, update output with data resource and html table
+							(<nb.IExecuteResult>this._outputs[i]).data = (<nb.IExecuteResult>output).data;
+							this._outputs[i].metadata = output.metadata;
+							added = true;
+							break;
+						}
+					}
+				}
+				break;
+			case 'execute_result_update':
+				let update = msg.content as nb.IExecuteResultUpdate;
+				// Send update to gridOutput component
+				this._onTableUpdated.fire({
+					resultSet: update.resultSet,
+					rows: update.data
+				});
+				break;
 			case 'display_data':
+				output = msg.content as nb.ICellOutput;
+				output.output_type = msgType;
+				// Display message outputs before grid outputs
+				if (this._outputs.length > 0) {
+					for (let i = 0; i < this._outputs.length; i++) {
+						if (this._outputs[i].output_type === 'execute_result') {
+							this._outputs.splice(i, 0, this.rewriteOutputUrls(output));
+							this.fireOutputsChanged();
+							added = true;
+							break;
+						}
+					}
+				}
+				break;
 			case 'stream':
 			case 'error':
 				output = msg.content as nb.ICellOutput;
@@ -698,25 +722,10 @@ export class CellModel extends Disposable implements ICellModel {
 		//     targets.push(model.length - 1);
 		//     this._displayIdMap.set(displayId, targets);
 		// }
-		if (output) {
+		if (output && !added) {
 			// deletes transient node in the serialized JSON
 			delete output['transient'];
-			// display message outputs before grid outputs
+			this._outputs.push(this.rewriteOutputUrls(output));
-			if (output.output_type === 'display_data' && this._outputs.length > 0) {
-				let added = false;
-				for (let i = 0; i < this._outputs.length; i++) {
-					if (this._outputs[i].output_type === 'execute_result') {
-						this._outputs.splice(i, 0, this.rewriteOutputUrls(output));
-						added = true;
-						break;
-					}
-				}
-				if (!added) {
-					this._outputs.push(this.rewriteOutputUrls(output));
-				}
-			} else {
-				this._outputs.push(this.rewriteOutputUrls(output));
-			}
 			// Only scroll on 1st output being added
 			let shouldScroll = this._outputs.length === 1;
 			this.fireOutputsChanged(shouldScroll);

src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts

@@ -21,6 +21,7 @@ import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilit
 import { NotebookModel } from 'sql/workbench/services/notebook/browser/models/notebookModel';
 import { IModelContentChangedEvent } from 'vs/editor/common/model/textModelEvents';
 import type { FutureInternal } from 'sql/workbench/services/notebook/browser/interfaces';
+import { ICellValue, ResultSetSummary } from 'sql/workbench/services/query/common/query';
 
 export interface ICellRange {
 	readonly start: number;
@@ -237,10 +238,6 @@ export interface INotebookModel {
 	 * Promise indicating when client session is ready to use.
 	 */
 	readonly sessionLoadFinished: Promise<void>;
-	/**
-	 * Promise indicating when output grid data is converted to mimeType and html.
-	 */
-	gridDataConversionComplete: Promise<any>;
 	/**
 	 * LanguageInfo saved in the notebook
 	 */
@@ -449,6 +446,11 @@ export interface IOutputChangedEvent {
 	shouldScroll: boolean;
 }
 
+export interface ITableUpdatedEvent {
+	resultSet: ResultSetSummary;
+	rows: ICellValue[][];
+}
+
 export interface ICellModel {
 	cellUri: URI;
 	id: string;
@@ -464,6 +466,7 @@ export interface ICellModel {
 	readonly outputs: ReadonlyArray<nb.ICellOutput>;
 	renderedOutputTextContent?: string[];
 	readonly onOutputsChanged: Event<IOutputChangedEvent>;
+	readonly onTableUpdated: Event<ITableUpdatedEvent>;
 	readonly onExecutionStateChange: Event<CellExecutionState>;
 	readonly executionState: CellExecutionState;
 	readonly notebookModel: NotebookModel;
@@ -491,9 +494,6 @@ export interface ICellModel {
 	readonly onCellMarkdownModeChanged: Event<boolean>;
 	sendChangeToNotebook(change: NotebookChangeType): void;
 	cellSourceChanged: boolean;
-	gridDataConversionComplete: Promise<void>;
-	addGridDataConversionPromise(complete: Promise<void>): void;
-	updateOutputData(batchId: number, id: number, data: any): void;
 }
 
 export interface IModelFactory {

src/sql/workbench/services/notebook/browser/models/notebookModel.ts

@@ -270,17 +270,6 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return this._sessionLoadFinished;
 	}
 
-	/**
-	 * Indicates all result grid output has been converted to mimeType and html.
-	 */
-	public get gridDataConversionComplete(): Promise<any> {
-		let promises = [];
-		for (let cell of this._cells) {
-			promises.push(cell.gridDataConversionComplete);
-		}
-		return Promise.all(promises);
-	}
-
 	/**
 	 * Notifies when the client session is ready for use
 	 */

src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts

@@ -6,7 +6,7 @@
 import { nb, IResultMessage } from 'azdata';
 import { localize } from 'vs/nls';
 import QueryRunner from 'sql/workbench/services/query/common/queryRunner';
-import { ResultSetSummary, ResultSetSubset, IColumn, BatchSummary } from 'sql/workbench/services/query/common/query';
+import { ResultSetSummary, IColumn, BatchSummary, ICellValue } from 'sql/workbench/services/query/common/query';
 import { IConnectionManagementService } from 'sql/platform/connection/common/connectionManagement';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import Severity from 'vs/base/common/severity';
@@ -368,6 +368,16 @@ class SqlKernel extends Disposable implements nb.IKernel {
 				}
 			}
 		}));
+		this._register(queryRunner.onResultSet(resultSet => {
+			if (this._future) {
+				this._future.handleResultSet(resultSet);
+			}
+		}));
+		this._register(queryRunner.onResultSetUpdate(resultSet => {
+			if (this._future) {
+				this._future.handleResultSetUpdate(resultSet);
+			}
+		}));
 		this._register(queryRunner.onBatchEnd(batch => {
 			if (this._future) {
 				this._future.handleBatchEnd(batch);
@@ -403,9 +413,14 @@ export class SQLFuture extends Disposable implements FutureInternal {
 	private doneDeferred = new Deferred<nb.IShellMessage>();
 	private configuredMaxRows: number = MAX_ROWS;
 	private _outputAddedPromises: Promise<void>[] = [];
-	private _querySubsetResultMap: Map<number, ResultSetSubset> = new Map<number, ResultSetSubset>();
 	private _errorOccurred: boolean = false;
 	private _stopOnError: boolean = true;
+	private _lastRowCountMap: Map<string, number> = new Map<string, number>();
+	// Map containing data resource and html table to be saved in notebook
+	private _dataToSaveMap: Map<string, any> = new Map<string, any>();
+	// Map containing row data returned from SQL Tools Service and used for table rendering
+	private _rowsMap: Map<string, any> = new Map<string, any>();
+
 	constructor(
 		private _queryRunner: QueryRunner | undefined,
 		private _executionCount: number | undefined,
@@ -463,7 +478,6 @@ export class SQLFuture extends Disposable implements FutureInternal {
 			this.doneHandler.handle(msg);
 		}
 		this.doneDeferred.resolve(msg);
-		this._querySubsetResultMap.clear();
 	}
 
 	sendInputReply(content: nb.IInputReply): void {
@@ -496,60 +510,81 @@ export class SQLFuture extends Disposable implements FutureInternal {
 		}
 	}
 
-	public handleBatchEnd(batch: BatchSummary): void {
+	public handleResultSet(resultSet: ResultSetSummary | ResultSetSummary[]): void {
-		if (this.ioHandler) {
+		let resultSets: ResultSetSummary[];
-			this._outputAddedPromises.push(this.processResultSets(batch));
+		if (!Array.isArray(resultSet)) {
+			resultSets = [resultSet];
+		} else {
+			resultSets = resultSet?.splice(0);
+		}
+		for (let set of resultSets) {
+			let key = set.batchId + '-' + set.id;
+			this._lastRowCountMap.set(key, 0);
+			// Convert the headers to data resource and html and send to cell model
+			let data = {
+				'application/vnd.dataresource+json': this.convertHeaderToDataResource(set.columnInfo),
+				'text/html': this.convertHeaderToHtmlTable(set.columnInfo)
+			};
+			this._dataToSaveMap.set(key, data);
+			this._rowsMap.set(key, []);
+			this.sendIOPubMessage(data, set);
 		}
 	}
 
-	private async processResultSets(batch: BatchSummary): Promise<void> {
+	public handleResultSetUpdate(resultSet: ResultSetSummary | ResultSetSummary[]): void {
-		try {
+		let resultSets: ResultSetSummary[];
-			let queryRowsPromises: Promise<void>[] = [];
+		if (!Array.isArray(resultSet)) {
-			for (let resultSet of batch.resultSetSummaries) {
+			resultSets = [resultSet];
-				let rowCount = resultSet.rowCount > this.configuredMaxRows ? this.configuredMaxRows : resultSet.rowCount;
+		} else {
-				if (rowCount === this.configuredMaxRows) {
+			resultSets = resultSet?.splice(0);
-					this.handleMessage(localize('sqlMaxRowsDisplayed', "Displaying Top {0} rows.", rowCount));
+		}
+		for (let set of resultSets) {
+			if (set.rowCount > this.configuredMaxRows) {
+				set.rowCount = this.configuredMaxRows;
+			}
+			let key = set.batchId + '-' + set.id;
+			if (set.rowCount !== this._lastRowCountMap.get(key)) {
+				this._outputAddedPromises.push(this.queryAndConvertData(set, this._lastRowCountMap.get(key)));
+				this._lastRowCountMap.set(key, set.rowCount);
+			}
+		}
+	}
+
+	public handleBatchEnd(batch: BatchSummary): void {
+		if (this.ioHandler) {
+			for (let set of batch.resultSetSummaries) {
+				if (set.rowCount > this.configuredMaxRows) {
+					this.handleMessage(localize('sqlMaxRowsDisplayed', "Displaying Top {0} rows.", this.configuredMaxRows));
 				}
-				queryRowsPromises.push(this.getAllQueryRows(rowCount, resultSet));
-			}
-			// We want to display table in the same order
-			let i = 0;
-			for (let resultSet of batch.resultSetSummaries) {
-				await queryRowsPromises[i];
-				this.sendResultSetAsIOPub(resultSet);
-				i++;
 			}
+		}
+	}
+
+	public async queryAndConvertData(resultSet: ResultSetSummary, lastRowCount: number): Promise<void> {
+		try {
+			let key = resultSet.batchId + '-' + resultSet.id;
+			// Query for rows and send rows to cell model
+			let queryResult = await this._queryRunner.getQueryRows(lastRowCount, resultSet.rowCount - lastRowCount, resultSet.batchId, resultSet.id);
+			this.sendIOPubUpdateMessage(queryResult.rows, resultSet);
+			let rows = this._rowsMap.get(key);
+			this._rowsMap.set(key, rows.concat(queryResult.rows));
+
+			// Convert rows to data resource and html and send to cell model to be saved
+			let dataResourceRows = this.convertRowsToDataResource(queryResult.rows);
+			let saveData = this._dataToSaveMap.get(key);
+			saveData['application/vnd.dataresource+json'].data = saveData['application/vnd.dataresource+json'].data.concat(dataResourceRows);
+			let htmlRows = this.convertRowsToHtml(queryResult.rows, key);
+			// Last value in array is '</table>' so we want to add row data before that
+			saveData['text/html'].splice(saveData['text/html'].length - 1, 0, ...htmlRows);
+			this._dataToSaveMap.set(key, saveData);
+			this.sendIOPubMessage(saveData, resultSet);
 		} catch (err) {
 			// TODO should we output this somewhere else?
 			this.logService.error(`Error outputting result sets from Notebook query: ${err}`);
 		}
 	}
 
-	private async getAllQueryRows(rowCount: number, resultSet: ResultSetSummary): Promise<void> {
+	private sendIOPubMessage(data: any, resultSet: ResultSetSummary): void {
-		let deferred: Deferred<void> = new Deferred<void>();
-		if (rowCount > 0) {
-			this._queryRunner.getQueryRows(0, rowCount, resultSet.batchId, resultSet.id).then((result) => {
-				this._querySubsetResultMap.set(resultSet.id, result);
-				deferred.resolve();
-			}, (err) => {
-				this._querySubsetResultMap.set(resultSet.id, { rowCount: 0, rows: [] });
-				deferred.reject(err);
-			});
-		} else {
-			this._querySubsetResultMap.set(resultSet.id, { rowCount: 0, rows: [] });
-			deferred.resolve();
-		}
-		return deferred;
-	}
-
-	private sendResultSetAsIOPub(resultSet: ResultSetSummary): void {
-		if (this._querySubsetResultMap && this._querySubsetResultMap.get(resultSet.id)) {
-			let subsetResult = this._querySubsetResultMap.get(resultSet.id);
-			this.sendIOPubMessage(subsetResult, resultSet);
-		}
-	}
-
-	private sendIOPubMessage(subsetResult: ResultSetSubset, resultSet: ResultSetSummary): void {
 		let msg: nb.IIOPubMessage = {
 			channel: 'iopub',
 			type: 'iopub',
@@ -559,18 +594,35 @@ export class SQLFuture extends Disposable implements FutureInternal {
 			},
 			content: <nb.IExecuteResult>{
 				output_type: 'execute_result',
-				metadata: {},
+				metadata: {
+					resultSet: resultSet
+				},
 				execution_count: this._executionCount,
-				data: {
+				data: data
-					'application/vnd.dataresource+json': this.convertToDataResource(resultSet.columnInfo, subsetResult),
+			},
-					'text/html': this.convertToHtmlTable(resultSet.columnInfo, subsetResult)
+			metadata: undefined,
-				}
+			parent_header: undefined
+		};
+		this.ioHandler.handle(msg);
+	}
+
+	private sendIOPubUpdateMessage(rows: any, resultSet: ResultSetSummary): void {
+		let msg: nb.IIOPubMessage = {
+			channel: 'iopub',
+			type: 'iopub',
+			header: <nb.IHeader>{
+				msg_id: undefined,
+				msg_type: 'execute_result_update'
+			},
+			content: <nb.IExecuteResultUpdate>{
+				output_type: 'execute_result_update',
+				resultSet: resultSet,
+				data: rows
 			},
 			metadata: undefined,
 			parent_header: undefined
 		};
 		this.ioHandler?.handle(msg);
-		this._querySubsetResultMap.delete(resultSet.id);
 	}
 
 	setIOPubHandler(handler: nb.MessageHandler<nb.IIOPubMessage>): void {
@@ -584,48 +636,53 @@ export class SQLFuture extends Disposable implements FutureInternal {
 		// no-op
 	}
 
-	private convertToDataResource(columns: IColumn[], subsetResult: ResultSetSubset): IDataResource {
+	private convertHeaderToDataResource(columns: IColumn[]): IDataResource {
 		let columnsResources: IDataResourceSchema[] = [];
 		columns.forEach(column => {
 			columnsResources.push({ name: escape(column.columnName) });
 		});
-		let columnsFields: IDataResourceFields = { fields: undefined };
+		let columnsFields: IDataResourceFields = { fields: columnsResources };
-		columnsFields.fields = columnsResources;
 		return {
 			schema: columnsFields,
-			data: subsetResult.rows.map(row => {
+			data: []
-				let rowObject: { [key: string]: any; } = {};
-				row.forEach((val, index) => {
-					rowObject[index] = val.displayValue;
-				});
-				return rowObject;
-			})
 		};
 	}
 
-	private convertToHtmlTable(columns: IColumn[], d: ResultSetSubset): string[] {
+	private convertHeaderToHtmlTable(columns: IColumn[]): string[] {
-		// Adding 3 for <table>, column title rows, </table>
+		let htmlTable: string[] = new Array(3);
-		let htmlStringArr: string[] = new Array(d.rowCount + 3);
+		htmlTable[0] = '<table>';
-		htmlStringArr[0] = '<table>';
 		if (columns.length > 0) {
 			let columnHeaders = '<tr>';
 			for (let column of columns) {
 				columnHeaders += `<th>${escape(column.columnName)}</th>`;
 			}
 			columnHeaders += '</tr>';
-			htmlStringArr[1] = columnHeaders;
+			htmlTable[1] = columnHeaders;
 		}
-		let i = 2;
+		htmlTable[2] = '</table>';
-		for (const row of d.rows) {
+		return htmlTable;
+	}
+
+	private convertRowsToDataResource(rows: ICellValue[][]): any[] {
+		return rows.map(row => {
+			let rowObject: { [key: string]: any; } = {};
+			row.forEach((val, index) => {
+				rowObject[index] = val.displayValue;
+			});
+			return rowObject;
+		});
+	}
+
+	private convertRowsToHtml(rows: ICellValue[][], key: string): string[] {
+		let htmlStringArr = [];
+		for (const row of rows) {
 			let rowData = '<tr>';
-			for (let columnIndex = 0; columnIndex < columns.length; columnIndex++) {
+			for (let columnIndex = 0; columnIndex < row.length; columnIndex++) {
 				rowData += `<td>${escape(row[columnIndex].displayValue)}</td>`;
 			}
 			rowData += '</tr>';
-			htmlStringArr[i] = rowData;
+			htmlStringArr.push(rowData);
-			i++;
 		}
-		htmlStringArr[htmlStringArr.length - 1] = '</table>';
 		return htmlStringArr;
 	}