diff --git a/src/sql/base/common/gridRange.ts b/src/sql/base/common/gridRange.ts index a397158b9c..8d3a981077 100644 --- a/src/sql/base/common/gridRange.ts +++ b/src/sql/base/common/gridRange.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import { IGridPosition, GridPosition } from 'sql/base/common/gridPosition'; +import { IRange } from 'vs/base/common/range'; import { isNumber } from 'vs/base/common/types'; /** @@ -358,4 +359,69 @@ export class GridRange { public static spansMultipleLines(range: IGridRange): boolean { return range.endRow > range.startRow; } + + /** + * Create an instance of IGridRange from Slick.Range. + */ + public static fromSlickRange(range: Slick.Range): IGridRange { + return { + startRow: range.fromRow, + endRow: range.toRow, + startColumn: range.fromCell, + endColumn: range.toCell + }; + } + + /** + * Create a list IGridRange from a list of Slick.Range. + */ + public static fromSlickRanges(ranges: Slick.Range[]): IGridRange[] { + return ranges.map(r => GridRange.fromSlickRange(r)); + } + + /** + * Merge the ranges by row or column and return merged ranges + * @param ranges the ranges to be merged + * @param mergeRows whether to merge the rows or columns. + */ + private static mergeRanges(ranges: IGridRange[], mergeRows: boolean): IRange[] { + let sourceRanges: IRange[] = ranges.map(r => { + if (mergeRows) { + return { start: r.startRow, end: r.endRow }; + } else { + return { start: r.startColumn, end: r.endColumn }; + } + }); + const mergedRanges: IRange[] = []; + sourceRanges = sourceRanges.sort((s1, s2) => { return s1.start - s2.start; }); + sourceRanges.forEach(range => { + let merged = false; + for (let i = 0; i < mergedRanges.length; i++) { + const mergedRange = mergedRanges[i]; + if (range.start <= mergedRange.end) { + mergedRange.end = Math.max(range.end, mergedRange.end); + merged = true; + break; + } + } + if (!merged) { + mergedRanges.push(range); + } + }); + return mergedRanges; + } + + /** + * Gets the unique row ranges. + */ + public static getUniqueRows(ranges: IGridRange[]): IRange[] { + return GridRange.mergeRanges(ranges, true); + } + + /** + * Gets the unique column ranges. + */ + public static getUniqueColumns(ranges: IGridRange[]): IRange[] { + return GridRange.mergeRanges(ranges, false); + } } diff --git a/src/sql/platform/query/common/query.ts b/src/sql/platform/query/common/query.ts index d734bd1192..bb873ff387 100644 --- a/src/sql/platform/query/common/query.ts +++ b/src/sql/platform/query/common/query.ts @@ -32,6 +32,7 @@ export interface IQueryEditorConfiguration { readonly openAfterSave: boolean; readonly showActionBar: boolean; readonly preferProvidersCopyHandler: boolean; + readonly promptForLargeRowSelection: boolean; }, readonly messages: { readonly showBatchTime: boolean; diff --git a/src/sql/workbench/contrib/query/browser/gridPanel.ts b/src/sql/workbench/contrib/query/browser/gridPanel.ts index 6dc073c8b9..06ea0a19f5 100644 --- a/src/sql/workbench/contrib/query/browser/gridPanel.ts +++ b/src/sql/workbench/contrib/query/browser/gridPanel.ts @@ -33,11 +33,11 @@ import { ActionBar, ActionsOrientation } from 'vs/base/browser/ui/actionbar/acti import { isInDOM, Dimension } from 'vs/base/browser/dom'; import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation'; import { IEditorService } from 'vs/workbench/services/editor/common/editorService'; -import { IAction, Separator } from 'vs/base/common/actions'; +import { IAction, Separator, toAction } from 'vs/base/common/actions'; import { ILogService } from 'vs/platform/log/common/log'; import { localize } from 'vs/nls'; import { IGridDataProvider } from 'sql/workbench/services/query/common/gridDataProvider'; -import { CancellationToken } from 'vs/base/common/cancellation'; +import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation'; import { GridPanelState, GridTableState } from 'sql/workbench/common/editor/query/gridTableState'; import { IUntitledTextEditorService } from 'vs/workbench/services/untitled/common/untitledTextEditorService'; import { SaveFormat } from 'sql/workbench/services/query/common/resultSerializer'; @@ -48,7 +48,7 @@ import { Orientation } from 'vs/base/browser/ui/splitview/splitview'; import { IQueryModelService } from 'sql/workbench/services/query/common/queryModel'; import { FilterButtonWidth, HeaderFilter } from 'sql/base/browser/ui/table/plugins/headerFilter.plugin'; import { HybridDataProvider } from 'sql/base/browser/ui/table/hybridDataProvider'; -import { INotificationService } from 'vs/platform/notification/common/notification'; +import { INotificationHandle, INotificationService, Severity } from 'vs/platform/notification/common/notification'; import { alert, status } from 'vs/base/browser/ui/aria/aria'; import { IExecutionPlanService } from 'sql/workbench/services/executionPlan/common/interfaces'; import { ExecutionPlanInput } from 'sql/workbench/contrib/executionPlan/browser/executionPlanInput'; @@ -58,6 +58,8 @@ import { IAccessibilityService } from 'vs/platform/accessibility/common/accessib import { IQuickInputService } from 'vs/platform/quickinput/common/quickInput'; import { queryEditorNullBackground } from 'sql/platform/theme/common/colorRegistry'; import { IComponentContextService } from 'sql/workbench/services/componentContext/browser/componentContextService'; +import { GridRange } from 'sql/base/common/gridRange'; +import { onUnexpectedError } from 'vs/base/common/errors'; const ROW_HEIGHT = 29; const HEADER_HEIGHT = 26; @@ -373,6 +375,7 @@ export abstract class GridTableBase extends Disposable implements IView, IQue private filterPlugin: HeaderFilter; private isDisposed: boolean = false; private gridConfig: IResultGridConfiguration; + private selectionChangeHandlerTokenSource: CancellationTokenSource | undefined; private columns: Slick.Column[]; @@ -664,7 +667,7 @@ export abstract class GridTableBase extends Disposable implements IView, IQue if (this.state) { this.state.selection = this.selectionModel.getSelectedRanges(); } - await this.notifyTableSelectionChanged(); + await this.handleTableSelectionChange(); }); this.table.grid.onScroll.subscribe((e, data) => { @@ -763,7 +766,7 @@ export abstract class GridTableBase extends Disposable implements IView, IQue this._state = val; } - private async getRowData(start: number, length: number): Promise { + private async getRowData(start: number, length: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise { let subset; if (this.dataProvider.isDataInMemory) { // handle the scenario when the data is sorted/filtered, @@ -771,23 +774,102 @@ export abstract class GridTableBase extends Disposable implements IView, IQue const data = await this.dataProvider.getRangeAsync(start, length); subset = data.map(item => Object.keys(item).map(key => item[key])); } else { - subset = (await this.gridDataProvider.getRowData(start, length)).rows; + subset = (await this.gridDataProvider.getRowData(start, length, cancellationToken, onProgressCallback)).rows; } return subset; } - private async notifyTableSelectionChanged() { - const selectedCells = []; - for (const range of this.state.selection) { - const subset = await this.getRowData(range.fromRow, range.toRow - range.fromRow + 1); - subset.forEach(row => { - // start with range.fromCell -1 because we have row number column which is not available in the actual data - for (let i = range.fromCell - 1; i < range.toCell; i++) { - selectedCells.push(row[i]); - } - }); + private async handleTableSelectionChange(): Promise { + if (this.selectionChangeHandlerTokenSource) { + this.selectionChangeHandlerTokenSource.cancel(); + } + this.selectionChangeHandlerTokenSource = new CancellationTokenSource(); + await this.notifyTableSelectionChanged(this.selectionChangeHandlerTokenSource); + } + + private async notifyTableSelectionChanged(cancellationTokenSource: CancellationTokenSource): Promise { + const gridRanges = GridRange.fromSlickRanges(this.state.selection ?? []); + const rowRanges = GridRange.getUniqueRows(gridRanges); + const columnRanges = GridRange.getUniqueColumns(gridRanges); + const rowCount = rowRanges.map(range => range.end - range.start + 1).reduce((p, c) => p + c); + const runAction = async (proceed: boolean) => { + const selectedCells = []; + if (proceed && !cancellationTokenSource.token.isCancellationRequested) { + let notificationHandle: INotificationHandle = undefined; + const timeout = setTimeout(() => { + notificationHandle = this.notificationService.notify({ + message: localize('resultsGrid.loadingData', "Loading selected rows for calculation..."), + severity: Severity.Info, + progress: { + infinite: true + }, + actions: { + primary: [ + toAction({ + id: 'cancelLoadingCells', + label: localize('resultsGrid.cancel', "Cancel"), + run: () => { + cancellationTokenSource.cancel(); + notificationHandle.close(); + } + })] + } + }); + }, 1000); + this.queryModelService.notifyCellSelectionChanged([]); + let rowsInProcessedRanges = 0; + for (const range of rowRanges) { + if (cancellationTokenSource.token.isCancellationRequested) { + break; + } + const rows = await this.getRowData(range.start, range.end - range.start + 1, cancellationTokenSource.token, (availableRows: number) => { + notificationHandle?.updateMessage(localize('resultsGrid.loadingDataWithProgress', "Loading selected rows for calculation ({0}/{1})...", rowsInProcessedRanges + availableRows, rowCount)); + }); + rows.forEach((row, rowIndex) => { + columnRanges.forEach(cr => { + for (let i = cr.start; i <= cr.end; i++) { + if (this.state.selection.some(selection => selection.contains(rowIndex + range.start, i))) { + // need to reduce the column index by 1 because we have row number column which is not available in the actual data + selectedCells.push(row[i - 1]); + } + } + }); + }); + rowsInProcessedRanges += range.end - range.start + 1; + } + clearTimeout(timeout); + notificationHandle?.close(); + } + cancellationTokenSource.dispose(); + if (!cancellationTokenSource.token.isCancellationRequested) { + this.queryModelService.notifyCellSelectionChanged(selectedCells); + } + }; + const showPromptConfigValue = this.configurationService.getValue('queryEditor').results.promptForLargeRowSelection; + if (this.options.inMemoryDataCountThreshold && rowCount > this.options.inMemoryDataCountThreshold && showPromptConfigValue) { + this.notificationService.prompt(Severity.Warning, localize('resultsGrid.largeRowSelectionPrompt.', 'You have selected {0} rows, it might take a while to load the data and calculate the summary, do you want to continue?', rowCount), [ + { + label: localize('resultsGrid.confirmLargeRowSelection', "Yes"), + run: async () => { + await runAction(true); + } + }, { + label: localize('resultsGrid.cancelLargeRowSelection', "Cancel"), + run: async () => { + await runAction(false); + } + }, { + label: localize('resultsGrid.donotShowLargeRowSelectionPromptAgain', "Don't show again"), + run: async () => { + this.configurationService.updateValue('queryEditor.results.promptForLargeRowSelection', false).catch(e => onUnexpectedError(e)); + await runAction(true); + }, + isSecondary: true + } + ]); + } else { + await runAction(true); } - this.queryModelService.notifyCellSelectionChanged(selectedCells); } private async onTableClick(event: ITableMouseEvent) { diff --git a/src/sql/workbench/contrib/query/browser/query.contribution.ts b/src/sql/workbench/contrib/query/browser/query.contribution.ts index a0e0965bd6..0bfdfd5744 100644 --- a/src/sql/workbench/contrib/query/browser/query.contribution.ts +++ b/src/sql/workbench/contrib/query/browser/query.contribution.ts @@ -424,6 +424,11 @@ const queryEditorConfiguration: IConfigurationNode = { 'description': localize('queryEditor.results.showActionBar', "Whether to show the action bar in the query results view"), 'default': true }, + 'queryEditor.results.promptForLargeRowSelection': { + 'type': 'boolean', + 'default': true, + 'description': localize('queryEditor.results.promptForLargeRowSelection', "When cells are selected in the results grid, ADS will calculate the summary for them, This setting controls whether to show the a confirmation when the number of rows selected is larger than the value specified in the 'inMemoryDataProcessingThreshold' setting. The default value is true.") + }, 'queryEditor.messages.showBatchTime': { 'type': 'boolean', 'description': localize('queryEditor.messages.showBatchTime', "Should execution time be shown for individual batches"), diff --git a/src/sql/workbench/contrib/query/browser/statusBarItems.ts b/src/sql/workbench/contrib/query/browser/statusBarItems.ts index e76c2c5ec1..bbc82aeb84 100644 --- a/src/sql/workbench/contrib/query/browser/statusBarItems.ts +++ b/src/sql/workbench/contrib/query/browser/statusBarItems.ts @@ -316,10 +316,12 @@ export class QueryResultSelectionSummaryStatusBarContribution extends Disposable const nullCount = selectedCells.filter(cell => cell.isNull).length; let summaryText, tooltipText; if (numericValues.length >= 2) { - const sum = numericValues.reduce((previous, current, idx, array) => previous + current); + const sum = numericValues.reduce((previous, current) => previous + current); + const min = numericValues.reduce((previous, current) => Math.min(previous, current)); + const max = numericValues.reduce((previous, current) => Math.max(previous, current)); summaryText = localize('status.query.summaryText', "Average: {0} Count: {1} Sum: {2}", Number((sum / numericValues.length).toFixed(3)), selectedCells.length, sum); tooltipText = localize('status.query.summaryTooltip', "Average: {0} Count: {1} Distinct Count: {2} Max: {3} Min: {4} Null Count: {5} Sum: {6}", - Number((sum / numericValues.length).toFixed(3)), selectedCells.length, distinctValues.size, Math.max(...numericValues), Math.min(...numericValues), nullCount, sum); + Number((sum / numericValues.length).toFixed(3)), selectedCells.length, distinctValues.size, max, min, nullCount, sum); } else { summaryText = summaryText = localize('status.query.summaryTextNonNumeric', "Count: {0} Distinct Count: {1} Null Count: {2}", selectedCells.length, distinctValues.size, nullCount); } diff --git a/src/sql/workbench/services/query/common/gridDataProvider.ts b/src/sql/workbench/services/query/common/gridDataProvider.ts index 4218ea9b30..e7c90e90b0 100644 --- a/src/sql/workbench/services/query/common/gridDataProvider.ts +++ b/src/sql/workbench/services/query/common/gridDataProvider.ts @@ -11,6 +11,8 @@ import { INotificationHandle, INotificationService, Severity } from 'vs/platform import * as nls from 'vs/nls'; import { toAction } from 'vs/base/common/actions'; import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService'; +import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation'; +import { GridRange } from 'sql/base/common/gridRange'; export interface IGridDataProvider { @@ -19,7 +21,7 @@ export interface IGridDataProvider { * @param rowStart 0-indexed start row to retrieve data from * @param numberOfRows total number of rows of data to retrieve */ - getRowData(rowStart: number, numberOfRows: number): Thenable; + getRowData(rowStart: number, numberOfRows: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Thenable; /** * Sends a copy request to copy data to the clipboard @@ -51,38 +53,9 @@ export interface IGridDataProvider { serializeResults(format: SaveFormat, selection: Slick.Range[]): Thenable; } -interface Range { - start: number; - end: number; -} - -/** - * Merge the ranges and get the sorted ranges. - */ -function mergeRanges(ranges: Range[]): Range[] { - const mergedRanges: Range[] = []; - const orderedRanges = ranges.sort((s1, s2) => { return s1.start - s2.start; }); - orderedRanges.forEach(range => { - let merged = false; - for (let i = 0; i < mergedRanges.length; i++) { - const mergedRange = mergedRanges[i]; - if (range.start <= mergedRange.end) { - mergedRange.end = Math.max(range.end, mergedRange.end); - merged = true; - break; - } - } - if (!merged) { - mergedRanges.push(range); - } - }); - return mergedRanges; -} - -export async function executeCopyWithNotification(notificationService: INotificationService, selections: Slick.Range[], isCancelable: boolean, copyHandler: (notification: INotificationHandle, rowCount: number) => Promise, onCanceled?: () => void): Promise { - const rowRanges: Range[] = mergeRanges(selections.map(selection => { return { start: selection.fromRow, end: selection.toRow }; })); +export async function executeCopyWithNotification(notificationService: INotificationService, selections: Slick.Range[], copyHandler: (notification: INotificationHandle, rowCount: number) => Promise, cancellationTokenSource?: CancellationTokenSource): Promise { + const rowRanges = GridRange.getUniqueRows(GridRange.fromSlickRanges(selections)); const rowCount = rowRanges.map(range => range.end - range.start + 1).reduce((p, c) => p + c); - let isCanceled = false; const notificationHandle = notificationService.notify({ message: nls.localize('gridDataProvider.copying', "Copying..."), severity: Severity.Info, @@ -90,13 +63,12 @@ export async function executeCopyWithNotification(notificationService: INotifica infinite: true }, actions: { - primary: isCancelable ? [ + primary: cancellationTokenSource ? [ toAction({ id: 'cancelCopyResults', label: nls.localize('gridDataProvider.cancelCopyResults', "Cancel"), run: () => { - isCanceled = true; - onCanceled!(); + cancellationTokenSource.cancel(); notificationHandle.close(); } })] : [] @@ -104,7 +76,7 @@ export async function executeCopyWithNotification(notificationService: INotifica }); try { await copyHandler(notificationHandle, rowCount); - if (!isCanceled) { + if (cancellationTokenSource === undefined || !cancellationTokenSource.token.isCancellationRequested) { notificationHandle.progress.done(); notificationHandle.updateActions({ primary: [ @@ -124,16 +96,16 @@ export async function executeCopyWithNotification(notificationService: INotifica } export async function copySelectionToClipboard(clipboardService: IClipboardService, notificationService: INotificationService, provider: IGridDataProvider, selections: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider): Promise { - let isCanceled = false; - await executeCopyWithNotification(notificationService, selections, true, async (notificationHandle, rowCount) => { - const batchSize = 100; + const cancellationTokenSource = new CancellationTokenSource() + await executeCopyWithNotification(notificationService, selections, async (notificationHandle, rowCount) => { const eol = provider.getEolString(); const valueSeparator = '\t'; const shouldRemoveNewLines = provider.shouldRemoveNewLines(); - // Merge the selections to get the columns and rows. - const columnRanges: Range[] = mergeRanges(selections.map(selection => { return { start: selection.fromCell, end: selection.toCell }; })); - const rowRanges: Range[] = mergeRanges(selections.map(selection => { return { start: selection.fromRow, end: selection.toRow }; })); + // Merge the selections to get the unique columns and unique rows. + const gridRanges = GridRange.fromSlickRanges(selections); + const columnRanges = GridRange.getUniqueColumns(gridRanges); + const rowRanges = GridRange.getUniqueRows(gridRanges); let processedRows = 0; const getMessageText = (): string => { @@ -151,57 +123,41 @@ export async function copySelectionToClipboard(clipboardService: IClipboardServi resultString = Array.from(headers.values()).join(valueSeparator).concat(eol); } - const batchResult: string[] = []; + const rowValues: string[] = []; for (const range of rowRanges) { + let rows: ICellValue[][]; + let processedRowsSnapshot = processedRows; + const rangeLength = range.end - range.start + 1; if (tableView && tableView.isDataInMemory) { - const rangeLength = range.end - range.start + 1; // If the data is sorted/filtered in memory, we need to get the data that is currently being displayed const tableData = await tableView.getRangeAsync(range.start, rangeLength); - const rowSet = tableData.map(item => Object.keys(item).map(key => item[key])); - batchResult.push(getStringValueForRowSet(rowSet, columnRanges, selections, range.start, eol, valueSeparator, shouldRemoveNewLines)); + rows = tableData.map(item => Object.keys(item).map(key => item[key])); processedRows += rangeLength; notificationHandle.updateMessage(getMessageText()); } else { - let start = range.start; - do { - const end = Math.min(start + batchSize - 1, range.end); - const batchLength = end - start + 1 - const rowSet = (await provider.getRowData(start, batchLength)).rows; - batchResult.push(getStringValueForRowSet(rowSet, columnRanges, selections, range.start, eol, valueSeparator, shouldRemoveNewLines)); - start = end + 1; - processedRows = processedRows + batchLength; - if (!isCanceled) { - notificationHandle.updateMessage(getMessageText()); - } - } while (start < range.end && !isCanceled) + rows = (await provider.getRowData(range.start, rangeLength, cancellationTokenSource.token, (fetchedRows) => { + processedRows = processedRowsSnapshot + fetchedRows; + notificationHandle.updateMessage(getMessageText()); + })).rows; } + rows.forEach((values, index) => { + const rowIndex = index + range.start; + const columnValues = []; + columnRanges.forEach(cr => { + for (let i = cr.start; i <= cr.end; i++) { + if (selections.some(selection => selection.contains(rowIndex, i))) { + columnValues.push(shouldRemoveNewLines ? removeNewLines(values[i].displayValue) : values[i].displayValue); + } + } + }); + rowValues.push(columnValues.join(valueSeparator)); + }); } - if (!isCanceled) { - resultString += batchResult.join(eol); + if (!cancellationTokenSource.token.isCancellationRequested) { + resultString += rowValues.join(eol); await clipboardService.writeText(resultString); } - }, () => { - isCanceled = true; - }); -} - -function getStringValueForRowSet(rows: ICellValue[][], columnRanges: Range[], selections: Slick.Range[], rowSetStartIndex: number, eol: string, valueSeparator: string, shouldRemoveNewLines: boolean): string { - let rowStrings: string[] = []; - rows.forEach((values, index) => { - const rowIndex = index + rowSetStartIndex; - const rowValues = []; - columnRanges.forEach(cr => { - for (let i = cr.start; i <= cr.end; i++) { - if (selections.some(selection => selection.contains(rowIndex, i))) { - rowValues.push(shouldRemoveNewLines ? removeNewLines(values[i].displayValue) : values[i].displayValue); - } else { - rowValues.push(''); - } - } - }); - rowStrings.push(rowValues.join(valueSeparator)); - }); - return rowStrings.join(eol); + }, cancellationTokenSource); } export function getTableHeaderString(provider: IGridDataProvider, selection: Slick.Range[]): string { diff --git a/src/sql/workbench/services/query/common/queryManagement.ts b/src/sql/workbench/services/query/common/queryManagement.ts index 31a6b77d77..f232e7ba09 100644 --- a/src/sql/workbench/services/query/common/queryManagement.ts +++ b/src/sql/workbench/services/query/common/queryManagement.ts @@ -18,6 +18,7 @@ import { ResultSetSubset } from 'sql/workbench/services/query/common/query'; import { isUndefined } from 'vs/base/common/types'; import { ILogService } from 'vs/platform/log/common/log'; import * as nls from 'vs/nls'; +import { CancellationToken } from 'vs/base/common/cancellation'; export const SERVICE_ID = 'queryManagementService'; @@ -50,7 +51,7 @@ export interface IQueryManagementService { runQueryString(ownerUri: string, queryString: string): Promise; runQueryAndReturn(ownerUri: string, queryString: string): Promise; parseSyntax(ownerUri: string, query: string): Promise; - getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise; + getQueryRows(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise; disposeQuery(ownerUri: string): Promise; changeConnectionUri(newUri: string, oldUri: string): Promise; saveResults(requestParams: azdata.SaveResultsRequestParams): Promise; @@ -271,9 +272,38 @@ export class QueryManagementService implements IQueryManagementService { }); } - public async getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise { - return this._runAction(rowData.ownerUri, (runner) => { - return runner.getQueryRows(rowData).then(r => r.resultSubset); + public async getQueryRows(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise { + const pageSize = 500; + return this._runAction(rowData.ownerUri, async (runner): Promise => { + const result = []; + let start = rowData.rowsStartIndex; + this._logService.trace(`Getting ${rowData.rowsCount} rows starting from index: ${rowData.rowsStartIndex}.`); + do { + const rowCount = Math.min(pageSize, rowData.rowsStartIndex + rowData.rowsCount - start); + this._logService.trace(`Paged Fetch - Getting ${rowCount} rows starting from index: ${start}.`); + const rowSet = await runner.getQueryRows({ + ownerUri: rowData.ownerUri, + batchIndex: rowData.batchIndex, + resultSetIndex: rowData.resultSetIndex, + rowsCount: rowCount, + rowsStartIndex: start + }); + this._logService.trace(`Paged Fetch - Received ${rowSet.resultSubset.rows} rows starting from index: ${start}.`); + result.push(...rowSet.resultSubset.rows); + start += rowCount; + if (onProgressCallback) { + onProgressCallback(start - rowData.rowsStartIndex); + } + } while (start < rowData.rowsStartIndex + rowData.rowsCount && (cancellationToken === undefined || !cancellationToken.isCancellationRequested)); + if (cancellationToken?.isCancellationRequested) { + this._logService.trace(`Stop getting more rows since cancellation has been requested.`); + } else { + this._logService.trace(`Successfully fetched ${result.length} rows. Expected Rows: ${rowData.rowsCount}.`); + } + return { + rows: result, + rowCount: result.length + }; }); } diff --git a/src/sql/workbench/services/query/common/queryRunner.ts b/src/sql/workbench/services/query/common/queryRunner.ts index ad0eb0b7da..5b6fe53854 100644 --- a/src/sql/workbench/services/query/common/queryRunner.ts +++ b/src/sql/workbench/services/query/common/queryRunner.ts @@ -30,6 +30,7 @@ import { IQueryEditorConfiguration } from 'sql/platform/query/common/query'; import { IDisposableDataProvider } from 'sql/base/common/dataProvider'; import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfiguration'; import { ICapabilitiesService } from 'sql/platform/capabilities/common/capabilitiesService'; +import { CancellationToken } from 'vs/base/common/cancellation'; /* * Query Runner class which handles running a query, reports the results to the content manager, @@ -419,7 +420,7 @@ export default class QueryRunner extends Disposable { /** * Get more data rows from the current resultSets from the service layer */ - public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number): Promise { + public getQueryRows(rowStart: number, numberOfRows: number, batchIndex: number, resultSetIndex: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise { let rowData: QueryExecuteSubsetParams = { ownerUri: this.uri, resultSetIndex: resultSetIndex, @@ -428,7 +429,7 @@ export default class QueryRunner extends Disposable { batchIndex: batchIndex }; - return this.queryManagementService.getQueryRows(rowData).then(r => r, error => { + return this.queryManagementService.getQueryRows(rowData, cancellationToken, onProgressCallback).then(r => r, error => { // this._notificationService.notify({ // severity: Severity.Error, // message: nls.localize('query.gettingRowsFailedError', 'Something went wrong getting more rows: {0}', error) @@ -566,8 +567,8 @@ export class QueryGridDataProvider implements IGridDataProvider { ) { } - getRowData(rowStart: number, numberOfRows: number): Promise { - return this.queryRunner.getQueryRows(rowStart, numberOfRows, this.batchId, this.resultSetId); + getRowData(rowStart: number, numberOfRows: number, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise { + return this.queryRunner.getQueryRows(rowStart, numberOfRows, this.batchId, this.resultSetId, cancellationToken, onProgressCallback); } copyResults(selection: Slick.Range[], includeHeaders?: boolean, tableView?: IDisposableDataProvider): Promise { @@ -590,7 +591,7 @@ export class QueryGridDataProvider implements IGridDataProvider { } private async handleCopyRequestByProvider(selections: Slick.Range[], includeHeaders?: boolean): Promise { - executeCopyWithNotification(this._notificationService, selections, false, async () => { + executeCopyWithNotification(this._notificationService, selections, async () => { await this.queryRunner.copyResults(selections, this.batchId, this.resultSetId, this.shouldRemoveNewLines(), this.shouldIncludeHeaders(includeHeaders)); }); } diff --git a/src/sql/workbench/services/query/test/common/queryRunner.test.ts b/src/sql/workbench/services/query/test/common/queryRunner.test.ts index a5878aa55f..a9a98aab69 100644 --- a/src/sql/workbench/services/query/test/common/queryRunner.test.ts +++ b/src/sql/workbench/services/query/test/common/queryRunner.test.ts @@ -59,7 +59,7 @@ suite('Query Runner', () => { const getRowStub = sinon.stub().returns(Promise.resolve(rowResults)); (instantiationService as TestInstantiationService).stub(IQueryManagementService, 'getQueryRows', getRowStub); const resultReturn = await runner.getQueryRows(0, 100, 0, 0); - assert(getRowStub.calledWithExactly({ ownerUri: uri, batchIndex: 0, resultSetIndex: 0, rowsStartIndex: 0, rowsCount: 100 })); + assert(getRowStub.calledWithExactly({ ownerUri: uri, batchIndex: 0, resultSetIndex: 0, rowsStartIndex: 0, rowsCount: 100 }, undefined, undefined)); assert.deepStrictEqual(resultReturn, rowResults); // batch complete const batchComplete: CompleteBatchSummary = { ...batch, executionEnd: 'endstring', executionElapsed: 'elapsedstring' }; diff --git a/src/sql/workbench/services/query/test/common/testQueryManagementService.ts b/src/sql/workbench/services/query/test/common/testQueryManagementService.ts index cd0f253e6b..482ef64833 100644 --- a/src/sql/workbench/services/query/test/common/testQueryManagementService.ts +++ b/src/sql/workbench/services/query/test/common/testQueryManagementService.ts @@ -10,6 +10,7 @@ import QueryRunner from 'sql/workbench/services/query/common/queryRunner'; import * as azdata from 'azdata'; import { IRange } from 'vs/editor/common/core/range'; import { ResultSetSubset } from 'sql/workbench/services/query/common/query'; +import { CancellationToken } from 'vs/base/common/cancellation'; export class TestQueryManagementService implements IQueryManagementService { _serviceBrand: undefined; @@ -50,7 +51,7 @@ export class TestQueryManagementService implements IQueryManagementService { parseSyntax(ownerUri: string, query: string): Promise { throw new Error('Method not implemented.'); } - getQueryRows(rowData: azdata.QueryExecuteSubsetParams): Promise { + getQueryRows(rowData: azdata.QueryExecuteSubsetParams, cancellationToken?: CancellationToken, onProgressCallback?: (availableRows: number) => void): Promise { throw new Error('Method not implemented.'); } async disposeQuery(ownerUri: string): Promise {