From e3ec6bf9c56717a51455a31c9ee2c110b5eab744 Mon Sep 17 00:00:00 2001 From: Lucy Zhang Date: Thu, 10 Sep 2020 13:31:40 -0700 Subject: [PATCH] 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 | 6 + .../browser/cellViews/output.component.ts | 20 ++ .../cellViews/outputArea.component.html | 4 +- .../browser/cellViews/outputArea.component.ts | 1 + .../browser/models/fileNotebookInput.ts | 6 +- .../notebook/browser/models/notebookInput.ts | 16 +- .../browser/models/untitledNotebookInput.ts | 6 +- .../browser/outputs/gridOutput.component.ts | 238 +++++++++++++++--- .../notebook/browser/outputs/mimeRegistry.ts | 3 + .../test/browser/cellToolbarActions.test.ts | 2 +- .../browser/dataResourceDataProvider.test.ts | 13 +- .../test/browser/notebookEditor.test.ts | 12 +- .../test/browser/notebookInput.test.ts | 7 +- .../workbench/contrib/notebook/test/stubs.ts | 3 + .../services/notebook/browser/models/cell.ts | 39 ++- .../browser/models/modelInterfaces.ts | 7 + .../notebook/browser/models/notebookModel.ts | 11 + .../notebook/browser/sql/sqlSessionManager.ts | 130 ++++------ .../services/query/common/queryManagement.ts | 5 + .../test/common/testQueryManagementService.ts | 3 + 20 files changed, 400 insertions(+), 132 deletions(-) diff --git a/src/sql/azdata.proposed.d.ts b/src/sql/azdata.proposed.d.ts index 1775fa1af5..ace915b4f7 100644 --- a/src/sql/azdata.proposed.d.ts +++ b/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' diff --git a/src/sql/workbench/contrib/notebook/browser/cellViews/output.component.ts b/src/sql/workbench/contrib/notebook/browser/cellViews/output.component.ts index 29a2128324..1b2ce5b5bf 100644 --- a/src/sql/workbench/contrib/notebook/browser/cellViews/output.component.ts +++ b/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 = componentRef.location.nativeElement; diff --git a/src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.html b/src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.html index abb1c02977..9c244ad0a3 100644 --- a/src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.html +++ b/src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.html @@ -6,7 +6,7 @@ -->
- +
-
\ No newline at end of file + diff --git a/src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.ts b/src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.ts index 4da99c5f0c..2efe310d66 100644 --- a/src/sql/workbench/contrib/notebook/browser/cellViews/outputArea.component.ts +++ b/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); } diff --git a/src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts index 5aeec05f89..d054b3d759 100644 --- a/src/sql/workbench/contrib/notebook/browser/models/fileNotebookInput.ts +++ b/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 { diff --git a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts index 66348e602f..6d6ac447ba 100644 --- a/src/sql/workbench/contrib/notebook/browser/models/notebookInput.ts +++ b/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; @@ -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 { + 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); diff --git a/src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts b/src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts index 3e421191fc..ba626f0bf0 100644 --- a/src/sql/workbench/contrib/notebook/browser/models/untitledNotebookInput.ts +++ b/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 { diff --git a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts b/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts index c8119ce908..49949e72d6 100644 --- a/src/sql/workbench/contrib/notebook/browser/outputs/gridOutput.component.ts +++ b/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 { ICellValue, ResultSetSummary, ResultSetSubset } from 'sql/workbench/services/query/common/query'; +import QueryRunner, { getEolString, shouldIncludeHeaders, shouldRemoveNewLines } from 'sql/workbench/services/query/common/queryRunner'; +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 { values } from 'vs/base/common/collections'; -import { assign } from 'vs/base/common/objects'; +import { SaveFormat, ResultSerializer, SaveResultsResponse } from 'sql/workbench/services/query/common/resultSerializer'; 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 = 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 = 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 { - 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 { @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 { this.cellOutput.metadata.azdata_chartOptions = options; this.cellModel.sendChangeToNotebook(NotebookChangeType.CellMetadataUpdated); } + + public convertData(set: ResultSetSummary): Promise { + return this._gridDataProvider.convertAllData(set); + } } export class DataResourceDataProvider implements IGridDataProvider { - private rows: ICellValue[][]; - constructor(source: IDataResource, - private resultSet: ResultSetSummary, - private documentUri: string, + 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, @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] = ''; + let columnHeaders = ''; + for (let column of this._resultSet.columnInfo) { + columnHeaders += ``; + } + columnHeaders += ''; + htmlTable[1] = columnHeaders; + htmlTable[2] = '
${escape(column.columnName)}
'; + + 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 { - let rowEnd = rowStart + numberOfRows; - if (rowEnd > this.rows.length) { - rowEnd = this.rows.length; + public async convertAllData(result: ResultSetSummary): Promise { + // Querying 50 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); + } + } + + 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 { + 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 { @@ -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 { - let serializer = this._instantiationService.createInstance(ResultSerializer); - return serializer.handleSerialization(this.documentUri, format, (filePath) => this.doSerialize(serializer, filePath, format, selection)); + if (this._queryRunner) { + 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 { @@ -312,10 +457,10 @@ export class DataResourceDataProvider implements IGridDataProvider { return Promise.resolve(undefined); } // TODO implement selection support - let columns = this.resultSet.columnInfo; - let rowLength = this.rows.length; + let columns = this._resultSet.columnInfo; + 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 = ''; + for (let columnIndex = 0; columnIndex < row.length; columnIndex++) { + rowData += `${escape(row[columnIndex].displayValue)}`; + } + rowData += ''; + htmlStringArr.push(rowData); + } + return htmlStringArr; + } } diff --git a/src/sql/workbench/contrib/notebook/browser/outputs/mimeRegistry.ts b/src/sql/workbench/contrib/notebook/browser/outputs/mimeRegistry.ts index 2f235c961e..917ae0544c 100644 --- a/src/sql/workbench/contrib/notebook/browser/outputs/mimeRegistry.ts +++ b/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; } diff --git a/src/sql/workbench/contrib/notebook/test/browser/cellToolbarActions.test.ts b/src/sql/workbench/contrib/notebook/test/browser/cellToolbarActions.test.ts index 2ba6af5d17..7cc9d398f9 100644 --- a/src/sql/workbench/contrib/notebook/test/browser/cellToolbarActions.test.ts +++ b/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 { +export async function createandLoadNotebookModel(codeContent?: nb.INotebookContents): Promise { let defaultCodeContent: nb.INotebookContents = { cells: [{ cell_type: CellTypes.Code, diff --git a/src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts b/src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts index a92f4e8681..100042afee 100644 --- a/src/sql/workbench/contrib/notebook/test/browser/dataResourceDataProvider.test.ts +++ b/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 documentUri = 'untitled:Notebook-0'; + let cellModel = TypeMoq.Mock.ofType(CellModel); + 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 diff --git a/src/sql/workbench/contrib/notebook/test/browser/notebookEditor.test.ts b/src/sql/workbench/contrib/notebook/test/browser/notebookEditor.test.ts index fc275fdf5b..f492b22bdc 100644 --- a/src/sql/workbench/contrib/notebook/test/browser/notebookEditor.test.ts +++ b/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 = [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; 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: { notebookUri: untitledNotebookInput.notebookUri } }); notebookService.addNotebookEditor(testNotebookEditor); @@ -669,6 +672,7 @@ function setupServices(arg: { workbenchThemeService?: WorkbenchThemeService, ins const uninstallEvent = new Emitter(); const didUninstallEvent = new Emitter(); + const notificationService = TypeMoq.Mock.ofType(TestNotificationService, TypeMoq.MockBehavior.Loose); const instantiationService = arg.instantiationService ?? 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: { 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) { diff --git a/src/sql/workbench/contrib/notebook/test/browser/notebookInput.test.ts b/src/sql/workbench/contrib/notebook/test/browser/notebookInput.test.ts index 78d6094278..62ab1a99cd 100644 --- a/src/sql/workbench/contrib/notebook/test/browser/notebookInput.test.ts +++ b/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 { 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); diff --git a/src/sql/workbench/contrib/notebook/test/stubs.ts b/src/sql/workbench/contrib/notebook/test/stubs.ts index a2e26e1061..500db2f2e3 100644 --- a/src/sql/workbench/contrib/notebook/test/stubs.ts +++ b/src/sql/workbench/contrib/notebook/test/stubs.ts @@ -43,6 +43,9 @@ export class NotebookModelStub implements INotebookModel { get sessionLoadFinished(): Promise { throw new Error('method not implemented.'); } + get gridDataConversionComplete(): Promise { + throw new Error('method not implemented.'); + } get notebookManagers(): INotebookManager[] { throw new Error('method not implemented.'); } diff --git a/src/sql/workbench/services/notebook/browser/models/cell.ts b/src/sql/workbench/services/notebook/browser/models/cell.ts index 7f03f73279..9bbe8d37cb 100644 --- a/src/sql/workbench/services/notebook/browser/models/cell.ts +++ b/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(); private _isCommandExecutionSettingEnabled: boolean = false; + private _gridDataConversionComplete: Promise[] = []; 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 { 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' + && (this._outputs[i]).batchId === batchId + && (this._outputs[i]).id === id) { + (this._outputs[i]).data = data; + break; + } + } + } + + public get gridDataConversionComplete(): Promise { + return Promise.all(this._gridDataConversionComplete).then(); + } + + public addGridDataConversionPromise(complete: Promise): 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); diff --git a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts b/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts index 1f2f29c141..445e522a80 100644 --- a/src/sql/workbench/services/notebook/browser/models/modelInterfaces.ts +++ b/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; + /** + * Promise indicating when output grid data is converted to mimeType and html. + */ + gridDataConversionComplete: Promise; /** * LanguageInfo saved in the notebook */ @@ -480,6 +484,9 @@ export interface ICellModel { showPreview: boolean; readonly onCellPreviewChanged: Event; sendChangeToNotebook(change: NotebookChangeType): void; + gridDataConversionComplete: Promise; + addGridDataConversionPromise(complete: Promise): void; + updateOutputData(batchId: number, id: number, data: any): void; } export interface IModelFactory { diff --git a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts b/src/sql/workbench/services/notebook/browser/models/notebookModel.ts index e47001aaec..8aaa3f9dbd 100644 --- a/src/sql/workbench/services/notebook/browser/models/notebookModel.ts +++ b/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 { + 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 */ diff --git a/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts b/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts index 5b5ba12d4a..d9b94dd1b0 100644 --- a/src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts +++ b/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(); private configuredMaxRows: number = MAX_ROWS; private _outputAddedPromises: Promise[] = []; - private _querySubsetResultMap: Map = new Map(); 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 { - try { - let queryRowsPromises: Promise[] = []; - for (let resultSet of batch.resultSetSummaries) { - let rowCount = resultSet.rowCount > this.configuredMaxRows ? this.configuredMaxRows : resultSet.rowCount; - if (rowCount === this.configuredMaxRows) { - this.handleMessage(localize('sqlMaxRowsDisplayed', "Displaying Top {0} rows.", rowCount)); + public onBatchEnd(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++; + } + } + + private async sendInitialResultSets(resultSet: ResultSetSummary | ResultSetSummary[]): Promise { + try { + 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 { - let deferred: Deferred = new Deferred(); - 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 { + private sendIOPubMessage(resultSet: ResultSetSummary, conversionComplete?: boolean, subsetResult?: ResultSetSubset): 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), - 'text/html': this.convertToHtmlTable(resultSet.columnInfo, subsetResult) - } + 'application/vnd.dataresource+json': this.convertToDataResource(resultSet.columnInfo), + '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): 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 }; - columnsFields.fields = columnsResources; + let columnsFields: IDataResourceFields = { fields: columnsResources }; return { schema: columnsFields, - data: subsetResult.rows.map(row => { - let rowObject: { [key: string]: any; } = {}; - row.forEach((val, index) => { - rowObject[index] = val.displayValue; - }); - return rowObject; - }) + data: [] }; } - private convertToHtmlTable(columns: IColumn[], d: ResultSetSubset): string[] { - // Adding 3 for , column title rows,
- let htmlStringArr: string[] = new Array(d.rowCount + 3); - htmlStringArr[0] = ''; + private convertToHtmlTable(columns: IColumn[]): string[] { + let htmlTable: string[] = new Array(3); + htmlTable[0] = '
'; if (columns.length > 0) { let columnHeaders = ''; for (let column of columns) { columnHeaders += ``; } columnHeaders += ''; - htmlStringArr[1] = columnHeaders; + htmlTable[1] = columnHeaders; } - let i = 2; - for (const row of d.rows) { - let rowData = ''; - for (let columnIndex = 0; columnIndex < columns.length; columnIndex++) { - rowData += ``; - } - rowData += ''; - htmlStringArr[i] = rowData; - i++; - } - htmlStringArr[htmlStringArr.length - 1] = '
${escape(column.columnName)}
${escape(row[columnIndex].displayValue)}
'; - return htmlStringArr; + htmlTable[2] = ''; + return htmlTable; } private convertToDisplayMessage(msg: IResultMessage | string): nb.IIOPubMessage { diff --git a/src/sql/workbench/services/query/common/queryManagement.ts b/src/sql/workbench/services/query/common/queryManagement.ts index 35838900fb..46c12c4cb8 100644 --- a/src/sql/workbench/services/query/common/queryManagement.ts +++ b/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; runQuery(ownerUri: string, range?: IRange, runOptions?: ExecutionPlanOptions): Promise; @@ -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 diff --git a/src/sql/workbench/services/query/test/common/testQueryManagementService.ts b/src/sql/workbench/services/query/test/common/testQueryManagementService.ts index bb11c35e03..098b6a113b 100644 --- a/src/sql/workbench/services/query/test/common/testQueryManagementService.ts +++ b/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 { return { messages: undefined }; }