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

* 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 09:44:10 -07:00
committed by GitHub
parent 1ff815fe5a
commit 945e04ed92
5 changed files with 32 additions and 51 deletions

View File

@@ -41,6 +41,7 @@ import { IQueryManagementService } from 'sql/workbench/services/query/common/que
import { values } from 'vs/base/common/collections';
import { URI } from 'vs/base/common/uri';
import { assign } from 'vs/base/common/objects';
import { escape } from 'sql/base/common/strings';
@Component({
selector: GridOutputComponent.SELECTOR,
@@ -372,7 +373,7 @@ export class DataResourceDataProvider implements IGridDataProvider {
}
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.
let numRows = 100;
for (let i = 0; i < result.rowCount; i += 100) {

View File

@@ -778,7 +778,7 @@ suite('Cell Model', function (): void {
test('Execute returns error status', async function (): Promise<void> {
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>{
content: <nb.IExecuteReply>{
execution_count: 1,
@@ -796,7 +796,7 @@ suite('Cell Model', function (): void {
test('Execute returns abort status', async function (): Promise<void> {
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>{
content: <nb.IExecuteReply>{
execution_count: 1,
@@ -815,7 +815,7 @@ suite('Cell Model', function (): void {
test('Execute throws exception', async function (): Promise<void> {
let testMsg = 'Test message';
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);
});

View File

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

View File

@@ -23,7 +23,6 @@ import { isUndefinedOrNull } from 'vs/base/common/types';
import { ILanguageMagic } from 'sql/workbench/services/notebook/browser/notebookService';
import { ITextResourcePropertiesService } from 'vs/editor/common/services/textResourceConfigurationService';
import { URI } from 'vs/base/common/uri';
import { getUriPrefix, uriPrefixes } from 'sql/platform/connection/common/utils';
import { firstIndex } from 'vs/base/common/arrays';
import { startsWith } from 'vs/base/common/strings';
import { onUnexpectedError } from 'vs/base/common/errors';
@@ -159,7 +158,7 @@ export class SqlSession implements nb.ISession {
}
class SqlKernel extends Disposable implements nb.IKernel {
private _queryRunner: QueryRunner;
private _queryRunners: Map<string, QueryRunner> = new Map<string, QueryRunner>();
private _currentConnection: IConnectionProfile;
private _currentConnectionProfile: ConnectionProfile;
static kernelId: number = 0;
@@ -168,7 +167,6 @@ class SqlKernel extends Disposable implements nb.IKernel {
private _future: SQLFuture;
private _executionCount: number = 0;
private _magicToExecutorMap = new Map<string, ExternalScriptMagic>();
private _connectionPath: string;
constructor(private _path: string,
@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
@@ -182,7 +180,6 @@ class SqlKernel extends Disposable implements nb.IKernel {
) {
super();
this.initMagics();
this.setConnectionPath();
}
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 {
if (this._id === undefined) {
this._id = (SqlKernel.kernelId++).toString();
@@ -265,30 +239,33 @@ class SqlKernel extends Disposable implements nb.IKernel {
public set connection(conn: IConnectionProfile) {
this._currentConnection = conn;
this._currentConnectionProfile = new ConnectionProfile(this._capabilitiesService, this._currentConnection);
this._queryRunner = undefined;
this._queryRunners.clear();
}
getSpec(): Thenable<nb.IKernelSpec> {
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 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
if (this._future && !this._queryRunner.hasCompleted) {
this._queryRunner.cancelQuery().then(ok => undefined, error => this._errorMessageService.showDialog(Severity.Error, sqlKernelError, error));
if (this._future && !queryRunner.hasCompleted) {
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
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) {
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));
queryRunner = this._instantiationService.createInstance(QueryRunner, queryRunnerUri);
this._queryRunners.set(queryRunnerUri, queryRunner);
this.queryManagementService.registerRunner(queryRunner, queryRunnerUri);
this._connectionManagementService.connect(this._currentConnectionProfile, queryRunnerUri).then((result) => {
this.addQueryEventListeners(queryRunner);
queryRunner.runQuery(code).catch(err => onUnexpectedError(err));
}).catch(err => onUnexpectedError(err));
} else {
canRun = false;
@@ -298,7 +275,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
// TODO verify this is "canonical" behavior
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) {
// Complete early
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> {
// 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 {
@@ -372,15 +349,16 @@ class SqlKernel extends Disposable implements nb.IKernel {
}
public async disconnect(): Promise<void> {
if (this._connectionPath) {
if (this._connectionManagementService.isConnected(this._connectionPath)) {
this._queryRunners.forEach(async (queryRunner: QueryRunner, uri: string) => {
if (this._connectionManagementService.isConnected(uri)) {
try {
await this._connectionManagementService.disconnect(this._connectionPath);
await this._connectionManagementService.disconnect(uri);
} catch (err) {
this.logService.error(err);
}
}
}
});
return;
}
}