Fix issues with changing SQL notebook connection while a cell is executing (#19624)

src/sql/workbench/api/common/notebooks/vscodeExecuteProvider.ts

@@ -11,6 +11,7 @@ import { addExternalInteractiveKernelMetadata, convertToVSCodeNotebookCell } fro
 import { CellTypes } from 'sql/workbench/services/notebook/common/contracts';
 import { VSCodeNotebookDocument } from 'sql/workbench/api/common/notebooks/vscodeNotebookDocument';
 import { URI } from 'vs/base/common/uri';
+import { notebookMultipleRequestsError } from 'sql/workbench/common/constants';
 
 class VSCodeFuture implements azdata.nb.IFuture {
 	private _inProgress = true;
@@ -149,7 +150,7 @@ class VSCodeKernel implements azdata.nb.IKernel {
 
 	requestExecute(content: azdata.nb.IExecuteRequest, disposeOnDone?: boolean): azdata.nb.IFuture {
 		if (this._activeRequest) {
-			throw new Error(nls.localize('notebookMultipleRequestsError', "Cannot execute code cell. Another cell is currently being executed."));
+			throw new Error(notebookMultipleRequestsError);
 		}
 		let executePromise: Promise<void>;
 		if (this._controller.executeHandler) {

src/sql/workbench/common/constants.ts

@@ -6,6 +6,7 @@
 import { RawContextKey } from 'vs/platform/contextkey/common/contextkey';
 import * as glob from 'vs/base/common/glob';
 import { SearchSortOrder } from 'vs/workbench/services/search/common/search';
+import * as nls from 'vs/nls';
 
 export const FindInNotebooksActionId = 'workbench.action.findInNotebooks';
 export const FocusActiveEditorCommandId = 'notebookSearch.action.focusActiveEditor';
@@ -84,3 +85,5 @@ export const RESULTS_GRID_DEFAULTS = {
 	cellPadding: [5, 8, 4],
 	rowHeight: 24
 };
+
+export const notebookMultipleRequestsError = nls.localize('notebookMultipleRequestsError', "Cannot execute code cell. Another cell is currently being executed.");

src/sql/workbench/services/notebook/browser/sql/sqlSessionManager.ts

@@ -27,6 +27,7 @@ import { getUriPrefix, uriPrefixes } from 'sql/platform/connection/common/utils'
 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 { notebookMultipleRequestsError } from 'sql/workbench/common/constants';
 
 export const sqlKernelError: string = localize("sqlKernelError", "SQL kernel error");
 export const MAX_ROWS = 5000;
@@ -215,6 +216,7 @@ class SqlKernel extends Disposable implements nb.IKernel {
 	private _executionCount: number = 0;
 	private _magicToExecutorMap = new Map<string, ExternalScriptMagic>();
 	private _connectionPath: string;
+	private _newConnection: boolean;
 
 	constructor(private _path: string,
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
@@ -310,7 +312,7 @@ 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._newConnection = true;
 	}
 
 	getSpec(): Thenable<nb.IKernelSpec> {
@@ -318,17 +320,17 @@ class SqlKernel extends Disposable implements nb.IKernel {
 	}
 
 	requestExecute(content: nb.IExecuteRequest, disposeOnDone?: boolean): nb.IFuture {
+		// Check if another cell is already running.
+		if (this._future?.inProgress) {
+			throw new Error(notebookMultipleRequestsError);
+		}
+
 		let canRun: boolean = true;
 		let code = this.getCodeWithoutCellMagic(content);
-		if (this._queryRunner) {
+		if (this._queryRunner && !this._newConnection) {
-			// Cancel any existing query
-			if (this._future && !this._queryRunner.hasCompleted) {
-				this._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));
 		} else if (this._currentConnection && this._currentConnectionProfile) {
+			this._newConnection = false;
 			this._queryRunner = this._instantiationService.createInstance(QueryRunner, this._connectionPath);
 			this.addQueryEventListeners(this._queryRunner);
 			this._connectionManagementService.connect(this._currentConnectionProfile, this._connectionPath).then((result) => {