ckaczor/azuredatastudio
1
0
Fork 0
mirror of https://github.com/ckaczor/azuredatastudio.git synced 2026-02-16 10:58:30 -05:00
Code Issues Packages Projects Releases Wiki Activity

Fix notebook table rendering with multiple code cells (#12363) (#12391)

Browse Source 
* create unique query runner for each cell

* use cellUri instead of cellId to identify runner

* disconnect each query runner connection

* remove queryrunners size check
This commit is contained in:
Lucy Zhang
2020-09-17 10:32:11 -07:00
committed by GitHub
parent 7b6328dccf
commit 0bc81e1078
5 changed files with 32 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/sql/azdata.d.ts vendored
View File

@@ -4902,6 +4902,8 @@ declare module 'azdata' {
* *
* @param disposeOnDone - Whether to dispose of the future when done. * @param disposeOnDone - Whether to dispose of the future when done.
* *
* @param cellId - Cell id (used by queryRunner)
*
* @returns A kernel future. * @returns A kernel future.
* *
* #### Notes * #### Notes
@@ -4916,7 +4918,7 @@ declare module 'azdata' {
* *
* **See also:** [[IExecuteReply]] * **See also:** [[IExecuteReply]]
*/ */
requestExecute(content: IExecuteRequest, disposeOnDone?: boolean): IFuture; requestExecute(content: IExecuteRequest, disposeOnDone?: boolean, cellUri?: string): IFuture;
/** /**
* Send a `complete_request` message. * Send a `complete_request` message.

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

@@ -41,6 +41,7 @@ import { IQueryManagementService } from 'sql/workbench/services/query/common/que
import { values } from 'vs/base/common/collections'; import { values } from 'vs/base/common/collections';
import { URI } from 'vs/base/common/uri'; import { URI } from 'vs/base/common/uri';
import { assign } from 'vs/base/common/objects'; import { assign } from 'vs/base/common/objects';
import { escape } from 'sql/base/common/strings';
@Component({ @Component({
selector: GridOutputComponent.SELECTOR, selector: GridOutputComponent.SELECTOR,
@@ -372,7 +373,7 @@ export class DataResourceDataProvider implements IGridDataProvider {
} }
public async convertAllData(result: ResultSetSummary): Promise<void> { public async convertAllData(result: ResultSetSummary): Promise<void> {
// Querying 50 rows at a time. Querying large amount of rows will be slow and // Querying 100 rows at a time. Querying large amount of rows will be slow and
// affect table rendering since each time the user scrolls, getRowData is called. // affect table rendering since each time the user scrolls, getRowData is called.
let numRows = 100; let numRows = 100;
for (let i = 0; i < result.rowCount; i += 100) { for (let i = 0; i < result.rowCount; i += 100) {

6
src/sql/workbench/contrib/notebook/test/electron-browser/cell.test.ts
View File

@@ -778,7 +778,7 @@ suite('Cell Model', function (): void {
test('Execute returns error status', async function (): Promise<void> { test('Execute returns error status', async function (): Promise<void> {
mockKernel.setup(k => k.requiresConnection).returns(() => false); mockKernel.setup(k => k.requiresConnection).returns(() => false);
mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
let replyMsg: nb.IExecuteReplyMsg = <nb.IExecuteReplyMsg>{ let replyMsg: nb.IExecuteReplyMsg = <nb.IExecuteReplyMsg>{
content: <nb.IExecuteReply>{ content: <nb.IExecuteReply>{
execution_count: 1, execution_count: 1,
@@ -796,7 +796,7 @@ suite('Cell Model', function (): void {
test('Execute returns abort status', async function (): Promise<void> { test('Execute returns abort status', async function (): Promise<void> {
mockKernel.setup(k => k.requiresConnection).returns(() => false); mockKernel.setup(k => k.requiresConnection).returns(() => false);
mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
let replyMsg: nb.IExecuteReplyMsg = <nb.IExecuteReplyMsg>{ let replyMsg: nb.IExecuteReplyMsg = <nb.IExecuteReplyMsg>{
content: <nb.IExecuteReply>{ content: <nb.IExecuteReply>{
execution_count: 1, execution_count: 1,
@@ -815,7 +815,7 @@ suite('Cell Model', function (): void {
test('Execute throws exception', async function (): Promise<void> { test('Execute throws exception', async function (): Promise<void> {
let testMsg = 'Test message'; let testMsg = 'Test message';
mockKernel.setup(k => k.requiresConnection).returns(() => false); mockKernel.setup(k => k.requiresConnection).returns(() => false);
mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => { mockKernel.setup(k => k.requestExecute(TypeMoq.It.isAny(), TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => {
throw new Error(testMsg); throw new Error(testMsg);
}); });

2
src/sql/workbench/services/notebook/browser/models/cell.ts
View File

@@ -408,7 +408,7 @@ export class CellModel extends Disposable implements ICellModel {
const future = kernel.requestExecute({ const future = kernel.requestExecute({
code: content, code: content,
stop_on_error: true stop_on_error: true
}, false); }, false, this._cellUri.toString());
this.setFuture(future as FutureInternal); this.setFuture(future as FutureInternal);
this.fireExecutionStateChanged(); this.fireExecutionStateChanged();
// For now, await future completion. Later we should just track and handle cancellation based on model notifications // For now, await future completion. Later we should just track and handle cancellation based on model notifications

68
src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts
View File

@@ -23,7 +23,6 @@ import { isUndefinedOrNull } from 'vs/base/common/types';
import { ILanguageMagic } from 'sql/workbench/services/notebook/browser/notebookService'; import { ILanguageMagic } from 'sql/workbench/services/notebook/browser/notebookService';
import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfigurationService'; import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfigurationService';
import { URI } from 'vs/base/common/uri'; import { URI } from 'vs/base/common/uri';
import { getUriPrefix, uriPrefixes } from 'sql/platform/connection/common/utils';
import { firstIndex } from 'vs/base/common/arrays'; import { firstIndex } from 'vs/base/common/arrays';
import { startsWith } from 'vs/base/common/strings'; import { startsWith } from 'vs/base/common/strings';
import { onUnexpectedError } from 'vs/base/common/errors'; import { onUnexpectedError } from 'vs/base/common/errors';
@@ -159,7 +158,7 @@ export class SqlSession implements nb.ISession {
} }
class SqlKernel extends Disposable implements nb.IKernel { class SqlKernel extends Disposable implements nb.IKernel {
private _queryRunner: QueryRunner; private _queryRunners: Map<string, QueryRunner> = new Map<string, QueryRunner>();
private _currentConnection: IConnectionProfile; private _currentConnection: IConnectionProfile;
private _currentConnectionProfile: ConnectionProfile; private _currentConnectionProfile: ConnectionProfile;
static kernelId: number = 0; static kernelId: number = 0;
@@ -168,7 +167,6 @@ class SqlKernel extends Disposable implements nb.IKernel {
private _future: SQLFuture; private _future: SQLFuture;
private _executionCount: number = 0; private _executionCount: number = 0;
private _magicToExecutorMap = new Map<string, ExternalScriptMagic>(); private _magicToExecutorMap = new Map<string, ExternalScriptMagic>();
private _connectionPath: string;
constructor(private _path: string, constructor(private _path: string,
@IConnectionManagementService private _connectionManagementService: IConnectionManagementService, @IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
@@ -182,7 +180,6 @@ class SqlKernel extends Disposable implements nb.IKernel {
) { ) {
super(); super();
this.initMagics(); this.initMagics();
this.setConnectionPath();
} }
private initMagics(): void { private initMagics(): void {
@@ -192,29 +189,6 @@ class SqlKernel extends Disposable implements nb.IKernel {
} }
} }
private setConnectionPath(): void {
if (this._path) {
let prefix = getUriPrefix(this._path);
if (!prefix || prefix === uriPrefixes.connection) {
this._connectionPath = uriPrefixes.notebook.concat(this._path);
} else if (prefix !== uriPrefixes.notebook) {
try {
let uri = URI.parse(this._path);
if (uri && uri.scheme) {
this._connectionPath = uri.toString().replace(uri.scheme, uriPrefixes.notebook);
}
} catch {
// Ignore exceptions from URI parsing
} finally {
// If _connectionPath hasn't been set yet, set _connectionPath to _path as a last resort
if (!this._connectionPath) {
this._connectionPath = this._path;
}
}
}
}
}
public get id(): string { public get id(): string {
if (this._id === undefined) { if (this._id === undefined) {
this._id = (SqlKernel.kernelId++).toString(); this._id = (SqlKernel.kernelId++).toString();
@@ -265,30 +239,33 @@ class SqlKernel extends Disposable implements nb.IKernel {
public set connection(conn: IConnectionProfile) { public set connection(conn: IConnectionProfile) {
this._currentConnection = conn; this._currentConnection = conn;
this._currentConnectionProfile = new ConnectionProfile(this._capabilitiesService, this._currentConnection); this._currentConnectionProfile = new ConnectionProfile(this._capabilitiesService, this._currentConnection);
this._queryRunner = undefined; this._queryRunners.clear();
} }
getSpec(): Thenable<nb.IKernelSpec> { getSpec(): Thenable<nb.IKernelSpec> {
return Promise.resolve(notebookConstants.sqlKernelSpec); return Promise.resolve(notebookConstants.sqlKernelSpec);
} }
requestExecute(content: nb.IExecuteRequest, disposeOnDone?: boolean): nb.IFuture { requestExecute(content: nb.IExecuteRequest, disposeOnDone?: boolean, cellUri?: string): nb.IFuture {
let canRun: boolean = true; let canRun: boolean = true;
let code = this.getCodeWithoutCellMagic(content); let code = this.getCodeWithoutCellMagic(content);
if (this._queryRunner) { let queryRunnerUri = 'queryRunner-' + cellUri;
let queryRunner: QueryRunner | undefined = this._queryRunners.get(queryRunnerUri);
if (queryRunner) {
// Cancel any existing query // Cancel any existing query
if (this._future && !this._queryRunner.hasCompleted) { if (this._future && !queryRunner.hasCompleted) {
this._queryRunner.cancelQuery().then(ok => undefined, error => this._errorMessageService.showDialog(Severity.Error, sqlKernelError, error)); queryRunner.cancelQuery().then(ok => undefined, error => this._errorMessageService.showDialog(Severity.Error, sqlKernelError, error));
// TODO when we can just show error as an output, should show an "execution canceled" error in output // TODO when we can just show error as an output, should show an "execution canceled" error in output
this._future.handleDone().catch(err => onUnexpectedError(err)); this._future.handleDone().catch(err => onUnexpectedError(err));
} }
this._queryRunner.runQuery(code).catch(err => onUnexpectedError(err)); queryRunner.runQuery(code).catch(err => onUnexpectedError(err));
} else if (this._currentConnection && this._currentConnectionProfile) { } else if (this._currentConnection && this._currentConnectionProfile) {
this._queryRunner = this._instantiationService.createInstance(QueryRunner, this._connectionPath); queryRunner = this._instantiationService.createInstance(QueryRunner, queryRunnerUri);
this.queryManagementService.registerRunner(this._queryRunner, this._connectionPath); this._queryRunners.set(queryRunnerUri, queryRunner);
this._connectionManagementService.connect(this._currentConnectionProfile, this._connectionPath).then((result) => { this.queryManagementService.registerRunner(queryRunner, queryRunnerUri);
this.addQueryEventListeners(this._queryRunner); this._connectionManagementService.connect(this._currentConnectionProfile, queryRunnerUri).then((result) => {
this._queryRunner.runQuery(code).catch(err => onUnexpectedError(err)); this.addQueryEventListeners(queryRunner);
queryRunner.runQuery(code).catch(err => onUnexpectedError(err));
}).catch(err => onUnexpectedError(err)); }).catch(err => onUnexpectedError(err));
} else { } else {
canRun = false; canRun = false;
@@ -298,7 +275,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
// TODO verify this is "canonical" behavior // TODO verify this is "canonical" behavior
let count = canRun ? ++this._executionCount : undefined; let count = canRun ? ++this._executionCount : undefined;
this._future = new SQLFuture(this._queryRunner, count, this._configurationService, this.logService); this._future = new SQLFuture(queryRunner, count, this._configurationService, this.logService);
if (!canRun) { if (!canRun) {
// Complete early // Complete early
this._future.handleDone(new Error(localize('connectionRequired', "A connection must be chosen to run notebook cells"))).catch(err => onUnexpectedError(err)); this._future.handleDone(new Error(localize('connectionRequired', "A connection must be chosen to run notebook cells"))).catch(err => onUnexpectedError(err));
@@ -334,8 +311,8 @@ class SqlKernel extends Disposable implements nb.IKernel {
interrupt(): Thenable<void> { interrupt(): Thenable<void> {
// TODO: figure out what to do with the QueryCancelResult // TODO: figure out what to do with the QueryCancelResult
return this._queryRunner.cancelQuery().then((cancelResult) => { let runners = [...this._queryRunners.values()];
}); return Promise.all(runners.map(queryRunner => queryRunner.cancelQuery())).then();
} }
private addQueryEventListeners(queryRunner: QueryRunner): void { private addQueryEventListeners(queryRunner: QueryRunner): void {
@@ -372,15 +349,16 @@ class SqlKernel extends Disposable implements nb.IKernel {
} }
public async disconnect(): Promise<void> { public async disconnect(): Promise<void> {
if (this._connectionPath) { this._queryRunners.forEach(async (queryRunner: QueryRunner, uri: string) => {
if (this._connectionManagementService.isConnected(this._connectionPath)) { if (this._connectionManagementService.isConnected(uri)) {
try { try {
await this._connectionManagementService.disconnect(this._connectionPath); await this._connectionManagementService.disconnect(uri);
} catch (err) { } catch (err) {
this.logService.error(err); this.logService.error(err);
} }
} }
} });
return; return;
} }
} }