Add grid streaming support for notebooks (#12175)

* add onResultUpdate handler in gridoutput

* convert rows to mimetype and html

* wait for data conversion to finish before saving

* detach changeRef after output is created

* fix save grid action

* move data conversion check to each cell

* move conversion logic to dataprovider

* notify data converting when user saves

* add comments and remove unused methods

* fix method return type

* fix tests

* fix convertData method header

* move azdata changes to azdata proposed

* address PR comments

* display top rows message

* fix messages/table ordering and query 100 rows

* add missing escape import

* set default max rows to 5000

* add undefined check to updateResultSet

* change gridDataConversionComplete return type

src/sql/azdata.proposed.d.ts

@@ -44,6 +44,12 @@ declare module 'azdata' {
 		export interface IKernelChangedArgs {
 			nbKernelAlias?: string
 		}
+
+		export interface IExecuteResult {
+			data: any;
+			batchId?: number;
+			id?: number;
+		}
 	}
 
 	export type SqlDbType = 'BigInt' | 'Binary' | 'Bit' | 'Char' | 'DateTime' | 'Decimal'

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

@@ -41,6 +41,9 @@ export class OutputComponent extends CellView implements OnInit, AfterViewInit {
 	private _initialized: boolean = false;
 	private _activeCellId: string;
 	private _componentInstance: IMimeComponent;
+	private _batchId?: number;
+	private _id?: number;
+	private _queryRunnerUri?: string;
 	public errorText: string;
 
 	constructor(
@@ -102,6 +105,18 @@ export class OutputComponent extends CellView implements OnInit, AfterViewInit {
 		return this._componentInstance;
 	}
 
+	@Input() set batchId(value: number) {
+		this._batchId = value;
+	}
+
+	@Input() set id(value: number) {
+		this._id = value;
+	}
+
+	@Input() set queryRunnerUri(value: string) {
+		this._queryRunnerUri = value;
+	}
+
 	get trustedMode(): boolean {
 		return this._trusted;
 	}
@@ -174,6 +189,11 @@ 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.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">
+		<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>
 	</div>
 </div>

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

@@ -38,6 +38,7 @@ 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,6 +10,7 @@ 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';
@@ -21,9 +22,10 @@ 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);
+		super(title, resource, textInput, textModelService, instantiationService, notebookService, extensionService, notificationService);
 	}
 
 	public get textInput(): FileEditorInput {

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

@@ -33,6 +33,9 @@ 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';
 
 export type ModeViewSaveHandler = (handle: number) => Thenable<boolean>;
 
@@ -221,7 +224,8 @@ 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 = [];
@@ -290,6 +294,16 @@ 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,6 +10,7 @@ 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';
@@ -21,9 +22,10 @@ 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);
+		super(title, resource, textInput, textModelService, instantiationService, notebookService, extensionService, notificationService);
 	}
 
 	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 } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
+import { IDataResource, MaxTableRowsConfigName, NotebookConfigSectionName, IDataResourceSchema, IDataResourceFields, MAX_ROWS } from 'sql/workbench/services/notebook/browser/sql/sqlSessionManager';
 import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfigurationService';
-import { getEolString, shouldIncludeHeaders, shouldRemoveNewLines } from 'sql/workbench/services/query/common/queryRunner';
+import QueryRunner, { getEolString, shouldIncludeHeaders, shouldRemoveNewLines } from 'sql/workbench/services/query/common/queryRunner';
-import { ICellValue, ResultSetSummary, ResultSetSubset } from 'sql/workbench/services/query/common/query';
+import { ResultSetSummary, ResultSetSubset, ICellValue, BatchSummary } 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';
@@ -30,16 +30,17 @@ import { GridTableBase } from 'sql/workbench/contrib/query/browser/gridPanel';
 import { getErrorMessage } from 'vs/base/common/errors';
 import { ISerializationService, SerializeDataParams } from 'sql/platform/serialization/common/serializationService';
 import { SaveResultAction, IGridActionContext } from 'sql/workbench/contrib/query/browser/actions';
-import { ResultSerializer, SaveResultsResponse, SaveFormat } from 'sql/workbench/services/query/common/resultSerializer';
+import { SaveFormat, ResultSerializer, SaveResultsResponse } from 'sql/workbench/services/query/common/resultSerializer';
-import { values } from 'vs/base/common/collections';
-import { assign } from 'vs/base/common/objects';
 import { IUntitledTextEditorService } from 'vs/workbench/services/untitled/common/untitledTextEditorService';
 import { ChartView } from 'sql/workbench/contrib/charts/browser/chartView';
 import { ToggleableAction } from 'sql/workbench/contrib/notebook/browser/notebookActions';
 import { IInsightOptions } from 'sql/workbench/common/editor/query/chartState';
 import { NotebookChangeType } from 'sql/workbench/services/notebook/common/contracts';
-import { URI } from 'vs/base/common/uri';
 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';
 
 @Component({
 	selector: GridOutputComponent.SELECTOR,
@@ -55,10 +56,17 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 	private _cellOutput: azdata.nb.ICellOutput;
 	private _bundleOptions: MimeModel.IOptions;
 	private _table: DataResourceTable;
+	private _batchId: number;
+	private _id: number;
+	private _queryRunnerUri: string;
+	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();
 	}
@@ -91,7 +99,28 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 		this._cellOutput = value;
 	}
 
+	@Input() set batchId(value: number) {
+		this._batchId = value;
+	}
+
+	@Input() set id(value: number) {
+		this._id = value;
+	}
+
+	@Input() set queryRunnerUri(value: string) {
+		this._queryRunnerUri = value;
+	}
+
 	ngOnInit() {
+		let config = this.configurationService.getValue(NotebookConfigSectionName);
+		if (config) {
+			let maxRows = config[MaxTableRowsConfigName] ? config[MaxTableRowsConfigName] : undefined;
+			if (maxRows && maxRows > 0) {
+				this._configuredMaxRows = maxRows;
+			}
+		}
+		// When a saved notebook is opened, there is no query runner
+		this._queryRunner = this.queryManagementService.getRunner(this._queryRunnerUri);
 		this.renderGrid();
 	}
 
@@ -102,16 +131,45 @@ 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, source, this.cellModel, this.cellOutput, state);
+			this._table = this.instantiationService.createInstance(DataResourceTable, this._batchId, this._id, this._queryRunner, 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);
@@ -122,11 +180,17 @@ export class GridOutputComponent extends AngularDisposable implements IMimeCompo
 
 class DataResourceTable extends GridTableBase<any> {
 
-	private _gridDataProvider: IGridDataProvider;
+	private _gridDataProvider: DataResourceDataProvider;
 	private _chart: ChartView;
 	private _chartContainer: HTMLElement;
+	private _batchId: number;
+	private _id: number;
+	private _queryRunner: QueryRunner;
 
-	constructor(source: IDataResource,
+	constructor(batchId: number,
+		id: number,
+		queryRunner: QueryRunner,
+		source: IDataResource,
 		private cellModel: ICellModel,
 		private cellOutput: azdata.nb.ICellOutput,
 		state: GridTableState,
@@ -137,7 +201,10 @@ class DataResourceTable extends GridTableBase<any> {
 		@IConfigurationService configurationService: IConfigurationService
 	) {
 		super(state, createResultSet(source), { actionOrientation: ActionsOrientation.HORIZONTAL }, contextMenuService, instantiationService, editorService, untitledEditorService, configurationService);
-		this._gridDataProvider = this.instantiationService.createInstance(DataResourceDataProvider, source, this.resultSet, this.cellModel.notebookModel.notebookUri.toString());
+		this._batchId = batchId;
+		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) {
@@ -223,25 +290,73 @@ class DataResourceTable extends GridTableBase<any> {
 		this.cellOutput.metadata.azdata_chartOptions = options;
 		this.cellModel.sendChangeToNotebook(NotebookChangeType.CellMetadataUpdated);
 	}
+
+	public convertData(set: ResultSetSummary): Promise<void> {
+		return this._gridDataProvider.convertAllData(set);
+	}
 }
 
 export class DataResourceDataProvider implements IGridDataProvider {
-	private rows: ICellValue[][];
+	private _rows: ICellValue[][];
-	constructor(source: IDataResource,
+	private _documentUri: string;
-		private resultSet: ResultSetSummary,
+	private _queryRunner: QueryRunner;
-		private documentUri: string,
+	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,
 		@INotificationService private _notificationService: INotificationService,
 		@IClipboardService private _clipboardService: IClipboardService,
 		@IConfigurationService private _configurationService: IConfigurationService,
 		@ITextResourcePropertiesService private _textResourcePropertiesService: ITextResourcePropertiesService,
 		@ISerializationService private _serializationService: ISerializationService,
-		@IInstantiationService private _instantiationService: IInstantiationService
+		@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 => {
+		this._rows = source.data.map(row => {
 			let rowData: azdata.DbCellValue[] = [];
 			Object.keys(row).forEach((val, index) => {
 				let displayValue = String(values(row)[index]);
@@ -256,16 +371,41 @@ export class DataResourceDataProvider implements IGridDataProvider {
 		});
 	}
 
-	getRowData(rowStart: number, numberOfRows: number): Thenable<ResultSetSubset> {
+	public async convertAllData(result: ResultSetSummary): Promise<void> {
-		let rowEnd = rowStart + numberOfRows;
+		// Querying 50 rows at a time. Querying large amount of rows will be slow and
-		if (rowEnd > this.rows.length) {
+		// affect table rendering since each time the user scrolls, getRowData is called.
-			rowEnd = this.rows.length;
+		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);
+		}
+	}
+
+	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) {
+			return this._queryRunner.getQueryRows(rowStart, numberOfRows, this._batchId, this._id);
+		} else {
+			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> {
@@ -282,7 +422,7 @@ export class DataResourceDataProvider implements IGridDataProvider {
 	}
 
 	getEolString(): string {
-		return getEolString(this._textResourcePropertiesService, this.documentUri);
+		return getEolString(this._textResourcePropertiesService, this._documentUri);
 	}
 	shouldIncludeHeaders(includeHeaders: boolean): boolean {
 		return shouldIncludeHeaders(includeHeaders, this._configurationService);
@@ -292,7 +432,7 @@ export class DataResourceDataProvider implements IGridDataProvider {
 	}
 
 	getColumnHeaders(range: Slick.Range): string[] {
-		let headers: string[] = this.resultSet.columnInfo.slice(range.fromCell, range.toCell + 1).map((info, i) => {
+		let headers: string[] = this._resultSet.columnInfo.slice(range.fromCell, range.toCell + 1).map((info, i) => {
 			return info.columnName;
 		});
 		return headers;
@@ -303,8 +443,13 @@ export class DataResourceDataProvider implements IGridDataProvider {
 	}
 
 	serializeResults(format: SaveFormat, selection: Slick.Range[]): Thenable<void> {
-		let serializer = this._instantiationService.createInstance(ResultSerializer);
+		if (this._queryRunner) {
-		return serializer.handleSerialization(this.documentUri, format, (filePath) => this.doSerialize(serializer, filePath, format, selection));
+			selection = selection ? selection : [new Slick.Range(0, 0, this._resultSet.rowCount - 1, this._resultSet.columnInfo.length - 1)];
+			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> {
@@ -312,10 +457,10 @@ export class DataResourceDataProvider implements IGridDataProvider {
 			return Promise.resolve(undefined);
 		}
 		// TODO implement selection support
-		let columns = this.resultSet.columnInfo;
+		let columns = this._resultSet.columnInfo;
-		let rowLength = this.rows.length;
+		let rowLength = this._rows.length;
 		let minRow = 0;
-		let maxRow = this.rows.length;
+		let maxRow = this._rows.length;
 		let singleSelection = selection && selection.length > 0 ? selection[0] : undefined;
 		if (singleSelection && this.isSelected(singleSelection)) {
 			rowLength = singleSelection.toRow - singleSelection.fromRow + 1;
@@ -345,7 +490,7 @@ export class DataResourceDataProvider implements IGridDataProvider {
 					return headerData;
 				}));
 			}
-			result = result.concat(this.rows.slice(index, endIndex).map(row => {
+			result = result.concat(this._rows.slice(index, endIndex).map(row => {
 				if (this.isSelected(singleSelection)) {
 					return row.slice(singleSelection.fromCell, singleSelection.toCell + 1);
 				} else {
@@ -371,6 +516,29 @@ 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

@@ -23,6 +23,9 @@ export interface IMimeComponent {
 	mimeType: string;
 	cellModel?: ICellModel;
 	cellOutput?: nb.ICellOutput;
+	batchId?: number;
+	id?: number;
+	queryRunnerUri?: string;
 	layout(): void;
 }
 

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

@@ -177,7 +177,7 @@ suite('CellToolbarActions', function (): void {
 	});
 });
 
-async function createandLoadNotebookModel(codeContent?: nb.INotebookContents): Promise<NotebookModel> {
+export async function createandLoadNotebookModel(codeContent?: nb.INotebookContents): Promise<NotebookModel> {
 	let defaultCodeContent: nb.INotebookContents = {
 		cells: [{
 			cell_type: CellTypes.Code,

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

@@ -21,6 +21,8 @@ import { SerializationService } from 'sql/platform/serialization/common/serializ
 import { SaveFormat, ResultSerializer } from 'sql/workbench/services/query/common/resultSerializer';
 import { InstantiationService } from 'vs/platform/instantiation/common/instantiationService';
 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';
 
 export class TestSerializationProvider implements azdata.SerializationProvider {
 	providerId: string;
@@ -63,8 +65,9 @@ suite('Data Resource Data Provider', function () {
 			id: 0,
 			rowCount: 2
 		};
-
+		let cellModel = TypeMoq.Mock.ofType(CellModel);
-		let documentUri = 'untitled:Notebook-0';
+		let notebookModel = await createandLoadNotebookModel();
+		cellModel.setup(x => x.notebookModel).returns(() => notebookModel);
 		tempFolderPath = path.join(os.tmpdir(), `TestDataResourceDataProvider_${uuid.v4()}`);
 		await fs.mkdir(tempFolderPath);
 
@@ -88,11 +91,13 @@ suite('Data Resource Data Provider', function () {
 		let _instantiationService = TypeMoq.Mock.ofType(InstantiationService, TypeMoq.MockBehavior.Strict);
 		_instantiationService.setup(x => x.createInstance(TypeMoq.It.isValue(ResultSerializer)))
 			.returns(() => serializer);
-
 		dataResourceDataProvider = new DataResourceDataProvider(
+			0, // batchId
+			0, // id
+			undefined, // QueryRunner
 			source,
 			resultSet,
-			documentUri,
+			cellModel.object,
 			_notificationService,
 			undefined, // IClipboardService
 			undefined, // IConfigurationService

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

@@ -57,6 +57,8 @@ 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 {
 	private _cells: Array<ICellModel> = [new CellModel(undefined, undefined)];
@@ -97,10 +99,11 @@ 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 } = setupServices({ instantiationService, workbenchThemeService }));
+		({ instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub, notificationService } = setupServices({ instantiationService, workbenchThemeService }));
 		// Create notebookEditor
 		notebookEditor = createNotebookEditor(instantiationService, workbenchThemeService, notebookService);
 	});
@@ -121,7 +124,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
+				undefined, instantiationService, notebookService, extensionService, notificationService.object
 			);
 			const testNotebookEditor = new NotebookEditorStub({ cellGuid: cellTextEditorGuid, editor: queryTextEditor, model: notebookModel, notebookParams: <INotebookParams>{ notebookUri: untitledNotebookInput.notebookUri } });
 			notebookService.addNotebookEditor(testNotebookEditor);
@@ -669,6 +672,7 @@ 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);
@@ -705,7 +709,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
+		undefined, instantiationService, notebookService, extensionService, notificationService.object
 	);
 
 	const cellTextEditorGuid = generateUuid();
@@ -720,7 +724,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 };
+	return { instantiationService, workbenchThemeService, notebookService, testTitle, extensionService, cellTextEditorGuid, queryTextEditor, untitledNotebookInput, notebookEditorStub, notificationService };
 }
 
 function createNotebookEditor(instantiationService: TestInstantiationService, workbenchThemeService: WorkbenchThemeService, notebookService: NotebookService) {

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

@@ -20,6 +20,7 @@ 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();
@@ -44,6 +45,8 @@ suite('Notebook Input', function (): void {
 
 	(instantiationService as TestInstantiationService).stub(INotebookService, mockNotebookService.object);
 
+	const mockNotificationService = TypeMoq.Mock.ofType(TestNotificationService);
+
 	let untitledTextInput: UntitledTextEditorInput;
 	let untitledNotebookInput: UntitledNotebookInput;
 
@@ -53,14 +56,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);
+			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object, mockNotificationService.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);
+			undefined, instantiationService, mockNotebookService.object, mockExtensionService.object, mockNotificationService.object);
 
 		let inputId = fileNotebookInput.getTypeId();
 		assert.strictEqual(inputId, FileNotebookInput.ID);

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

@@ -43,6 +43,9 @@ 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

@@ -64,6 +64,7 @@ export class CellModel extends Disposable implements ICellModel {
 	private _showPreview: boolean = true;
 	private _onCellPreviewChanged = new Emitter<boolean>();
 	private _isCommandExecutionSettingEnabled: boolean = false;
+	private _gridDataConversionComplete: Promise<void>[] = [];
 
 	constructor(cellData: nb.ICellContents,
 		private _options: ICellModelOptions,
@@ -338,6 +339,8 @@ 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;
@@ -523,6 +526,25 @@ 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;
 	}
@@ -579,7 +601,22 @@ export class CellModel extends Disposable implements ICellModel {
 		if (output) {
 			// deletes transient node in the serialized JSON
 			delete output['transient'];
-			this._outputs.push(this.rewriteOutputUrls(output));
+			// display message outputs before grid outputs
+			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

@@ -237,6 +237,10 @@ 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
 	 */
@@ -480,6 +484,9 @@ export interface ICellModel {
 	showPreview: boolean;
 	readonly onCellPreviewChanged: Event<boolean>;
 	sendChangeToNotebook(change: NotebookChangeType): void;
+	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

@@ -271,6 +271,17 @@ 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 { BatchSummary, ResultSetSummary, IColumn, ResultSetSubset } from 'sql/workbench/services/query/common/query';
+import { ResultSetSummary, ResultSetSubset, IColumn, BatchSummary } 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';
@@ -14,8 +14,8 @@ import { Deferred } from 'sql/base/common/promise';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { IErrorMessageService } from 'sql/platform/errorMessage/common/errorMessageService';
 import { ConnectionProfile } from 'sql/platform/connection/common/connectionProfile';
-import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { escape } from 'sql/base/common/strings';
+import { IConnectionProfile } from 'sql/platform/connection/common/interfaces';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService';
 import { ILogService } from 'vs/platform/log/common/log';
@@ -29,6 +29,7 @@ import { startsWith } from 'vs/base/common/strings';
 import { onUnexpectedError } from 'vs/base/common/errors';
 import { FutureInternal, notebookConstants } from 'sql/workbench/services/notebook/browser/interfaces';
 import { tryMatchCellMagic } from 'sql/workbench/services/notebook/browser/utils';
+import { IQueryManagementService } from 'sql/workbench/services/query/common/queryManagement';
 
 export const sqlKernelError: string = localize("sqlKernelError", "SQL kernel error");
 export const MAX_ROWS = 5000;
@@ -176,7 +177,8 @@ class SqlKernel extends Disposable implements nb.IKernel {
 		@IErrorMessageService private _errorMessageService: IErrorMessageService,
 		@IConfigurationService private _configurationService: IConfigurationService,
 		@ILogService private readonly logService: ILogService,
-		@ITextResourcePropertiesService private readonly textResourcePropertiesService: ITextResourcePropertiesService
+		@ITextResourcePropertiesService private readonly textResourcePropertiesService: ITextResourcePropertiesService,
+		@IQueryManagementService private queryManagementService: IQueryManagementService
 	) {
 		super();
 		this.initMagics();
@@ -283,6 +285,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 			this._queryRunner.runQuery(code).catch(err => onUnexpectedError(err));
 		} else if (this._currentConnection && this._currentConnectionProfile) {
 			this._queryRunner = this._instantiationService.createInstance(QueryRunner, this._connectionPath);
+			this.queryManagementService.registerRunner(this._queryRunner, this._connectionPath);
 			this._connectionManagementService.connect(this._currentConnectionProfile, this._connectionPath).then((result) => {
 				this.addQueryEventListeners(this._queryRunner);
 				this._queryRunner.runQuery(code).catch(err => onUnexpectedError(err));
@@ -349,9 +352,14 @@ class SqlKernel extends Disposable implements nb.IKernel {
 				}
 			}
 		}));
+		this._register(queryRunner.onResultSet(resultSet => {
+			if (this._future) {
+				this._future.onResultSet(resultSet);
+			}
+		}));
 		this._register(queryRunner.onBatchEnd(batch => {
 			if (this._future) {
-				this._future.handleBatchEnd(batch);
+				this._future.onBatchEnd(batch);
 			}
 		}));
 	}
@@ -384,9 +392,9 @@ 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;
+
 	constructor(
 		private _queryRunner: QueryRunner,
 		private _executionCount: number | undefined,
@@ -442,7 +450,6 @@ export class SQLFuture extends Disposable implements FutureInternal {
 			this.doneHandler.handle(msg);
 		}
 		this.doneDeferred.resolve(msg);
-		this._querySubsetResultMap.clear();
 	}
 
 	sendInputReply(content: nb.IInputReply): void {
@@ -473,28 +480,32 @@ export class SQLFuture extends Disposable implements FutureInternal {
 		}
 	}
 
-	public handleBatchEnd(batch: BatchSummary): void {
+	public onResultSet(resultSet: ResultSetSummary | ResultSetSummary[]): void {
 		if (this.ioHandler) {
-			this._outputAddedPromises.push(this.processResultSets(batch));
+			this._outputAddedPromises.push(this.sendInitialResultSets(resultSet));
 		}
 	}
 
-	private async processResultSets(batch: BatchSummary): Promise<void> {
+	public onBatchEnd(batch: BatchSummary): void {
-		try {
+		if (this.ioHandler) {
-			let queryRowsPromises: Promise<void>[] = [];
+			for (let set of batch.resultSetSummaries) {
-			for (let resultSet of batch.resultSetSummaries) {
+				if (set.rowCount > this.configuredMaxRows) {
-				let rowCount = resultSet.rowCount > this.configuredMaxRows ? this.configuredMaxRows : resultSet.rowCount;
+					this.handleMessage(localize('sqlMaxRowsDisplayed', "Displaying Top {0} rows.", this.configuredMaxRows));
-				if (rowCount === this.configuredMaxRows) {
-					this.handleMessage(localize('sqlMaxRowsDisplayed', "Displaying Top {0} rows.", rowCount));
 				}
-				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];
+	private async sendInitialResultSets(resultSet: ResultSetSummary | ResultSetSummary[]): Promise<void> {
-				this.sendResultSetAsIOPub(resultSet);
+		try {
-				i++;
+			let resultsToAdd: ResultSetSummary[];
+			if (!Array.isArray(resultSet)) {
+				resultsToAdd = [resultSet];
+			} else {
+				resultsToAdd = resultSet?.splice(0);
+			}
+			for (let set of resultsToAdd) {
+				this.sendIOPubMessage(set, false);
 			}
 		} catch (err) {
 			// TODO should we output this somewhere else?
@@ -502,31 +513,7 @@ export class SQLFuture extends Disposable implements FutureInternal {
 		}
 	}
 
-	private async getAllQueryRows(rowCount: number, resultSet: ResultSetSummary): Promise<void> {
+	private sendIOPubMessage(resultSet: ResultSetSummary, conversionComplete?: boolean, subsetResult?: ResultSetSubset): 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',
@@ -538,16 +525,21 @@ export class SQLFuture extends Disposable implements FutureInternal {
 				output_type: 'execute_result',
 				metadata: {},
 				execution_count: this._executionCount,
+				// Initial data sent to notebook only contains column headers since
+				// onResultSet only returns the column info (and no row data).
+				// Row data conversion will be handled in DataResourceDataProvider
 				data: {
-					'application/vnd.dataresource+json': this.convertToDataResource(resultSet.columnInfo, subsetResult),
+					'application/vnd.dataresource+json': this.convertToDataResource(resultSet.columnInfo),
-					'text/html': this.convertToHtmlTable(resultSet.columnInfo, subsetResult)
+					'text/html': this.convertToHtmlTable(resultSet.columnInfo)
-				}
+				},
+				batchId: resultSet.batchId,
+				id: resultSet.id,
+				queryRunnerUri: this._queryRunner.uri,
 			},
 			metadata: undefined,
 			parent_header: undefined
 		};
 		this.ioHandler.handle(msg);
-		this._querySubsetResultMap.delete(resultSet.id);
 	}
 
 	setIOPubHandler(handler: nb.MessageHandler<nb.IIOPubMessage>): void {
@@ -561,49 +553,31 @@ export class SQLFuture extends Disposable implements FutureInternal {
 		// no-op
 	}
 
-	private convertToDataResource(columns: IColumn[], subsetResult: ResultSetSubset): IDataResource {
+	private convertToDataResource(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 convertToHtmlTable(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;
-			let rowData = '<tr>';
-			for (let columnIndex = 0; columnIndex < columns.length; columnIndex++) {
-				rowData += `<td>${escape(row[columnIndex].displayValue)}</td>`;
-			}
-			rowData += '</tr>';
-			htmlStringArr[i] = rowData;
-			i++;
-		}
-		htmlStringArr[htmlStringArr.length - 1] = '</table>';
-		return htmlStringArr;
 	}
 
 	private convertToDisplayMessage(msg: IResultMessage | string): nb.IIOPubMessage {

src/sql/workbench/services/query/common/queryManagement.ts

@@ -39,6 +39,7 @@ export interface IQueryManagementService {
 	isProviderRegistered(providerId: string): boolean;
 	getRegisteredProviders(): string[];
 	registerRunner(runner: QueryRunner, uri: string): void;
+	getRunner(uri: string): QueryRunner | undefined;
 
 	cancelQuery(ownerUri: string): Promise<QueryCancelResult>;
 	runQuery(ownerUri: string, range?: IRange, runOptions?: ExecutionPlanOptions): Promise<void>;
@@ -138,6 +139,10 @@ export class QueryManagementService implements IQueryManagementService {
 		}
 	}
 
+	public getRunner(uri: string): QueryRunner | undefined {
+		return this._queryRunners.get(uri);
+	}
+
 	// Handles logic to run the given handlerCallback at the appropriate time. If the given runner is
 	// undefined, the handlerCallback is put on the _handlerCallbackQueue to be run once the runner is set
 	// public for testing only

src/sql/workbench/services/query/test/common/testQueryManagementService.ts

@@ -26,6 +26,9 @@ export class TestQueryManagementService implements IQueryManagementService {
 	registerRunner(runner: QueryRunner, uri: string): void {
 		return;
 	}
+	getRunner(uri: string): QueryRunner {
+		throw new Error('Method not implemented.');
+	}
 	async cancelQuery(ownerUri: string): Promise<azdata.QueryCancelResult> {
 		return { messages: undefined };
 	}