Transfer URI Correctly (#24614)

* transfer uri regardless of queryrunner presence

* throw error if new URI already exists

* bump sts version

extensions/mssql/config.json

@@ -1,6 +1,6 @@
 {
 	"downloadUrl": "https://github.com/Microsoft/sqltoolsservice/releases/download/{#version#}/microsoft.sqltools.servicelayer-{#fileName#}",
-	"version": "4.10.0.14",
+	"version": "4.10.0.15",
 	"downloadFileNames": {
 		"Windows_86": "win-x86-net7.0.zip",
 		"Windows_64": "win-x64-net7.0.zip",

src/sql/workbench/services/query/common/queryManagement.ts

@@ -339,16 +339,16 @@ export class QueryManagementService implements IQueryManagementService {
 		let item = this._queryRunners.get(oldUri);
 		if (!item) {
 			this._logService.error(`No query runner found for old URI : '${oldUri}'`);
-			throw new Error(nls.localize('queryManagement.noQueryRunnerForUri', 'Could not find Query Runner for uri: {0}', oldUri));
-		}
-		if (this._queryRunners.get(newUri)) {
+		} else if (this._queryRunners.get(newUri)) {
 			this._logService.error(`New URI : '${newUri}' already has a query runner.`);
 			throw new Error(nls.localize('queryManagement.uriAlreadyHasQueryRunner', 'Uri: {0} unexpectedly already has a query runner.', newUri));
+		} else {
+			this._queryRunners.set(newUri, item);
+			this._queryRunners.delete(oldUri);
 		}
-		this._queryRunners.set(newUri, item);
-		this._queryRunners.delete(oldUri);
-		return this._runAction(newUri, (runner) => {
-			return runner.connectionUriChanged(newUri, oldUri);
+
+		return this._runAction(newUri, (handler) => {
+			return handler.connectionUriChanged(newUri, oldUri);
 		});
 	}
 

src/sql/workbench/services/query/common/queryModelService.ts

@@ -22,6 +22,7 @@ import Severity from 'vs/base/common/severity';
 import EditQueryRunner from 'sql/workbench/services/editData/common/editQueryRunner';
 import { IRange } from 'vs/editor/common/core/range';
 import { ClipboardData, IClipboardService } from 'vs/platform/clipboard/common/clipboardService';
+import { IQueryManagementService } from 'sql/workbench/services/query/common/queryManagement';
 
 const selectionSnippetMaxLen = 100;
 
@@ -81,6 +82,7 @@ export class QueryModelService implements IQueryModelService {
 
 	// CONSTRUCTOR /////////////////////////////////////////////////////////
 	constructor(
+		@IQueryManagementService protected queryManagementService: IQueryManagementService,
 		@IInstantiationService private _instantiationService: IInstantiationService,
 		@INotificationService private _notificationService: INotificationService,
 		@IClipboardService private _clipboardService: IClipboardService,
@@ -437,21 +439,17 @@ export class QueryModelService implements IQueryModelService {
 	}
 
 	public async changeConnectionUri(newUri: string, oldUri: string): Promise<void> {
-		// Get existing query runner
-		let queryRunner = this.internalGetQueryRunner(oldUri);
-		if (!queryRunner) {
-			// Nothing to do if we don't have a query runner currently (no connection)
-			return;
-		}
-		else if (this._queryInfoMap.has(newUri)) {
+		if (this._queryInfoMap.has(newUri)) {
 			this._logService.error(`New URI '${newUri}' already has query info associated with it.`);
 			throw new Error(nls.localize('queryModelService.uriAlreadyHasQuery', '{0} already has an existing query', newUri));
 		}
+		await this.queryManagementService.changeConnectionUri(newUri, oldUri);
 
-		await queryRunner.changeConnectionUri(newUri, oldUri);
-
-		// remove the old key and set new key with same query info as old uri. (Info existence is checked in internalGetQueryRunner)
+		// remove the old key and set new key with same query info as old uri.
 		let info = this._queryInfoMap.get(oldUri);
+		if (!info) {
+			return;
+		}
 		info.uri = newUri;
 		this._queryInfoMap.set(newUri, info);
 		this._queryInfoMap.delete(oldUri);

src/sql/workbench/services/query/common/queryRunner.ts

@@ -468,10 +468,6 @@ export default class QueryRunner extends Disposable {
 		this._isDisposed = true;
 	}
 
-	public changeConnectionUri(oldUri: string, newUri: string): Promise<void> {
-		return this.queryManagementService.changeConnectionUri(oldUri, newUri);
-	}
-
 	get totalElapsedMilliseconds(): number {
 		return this._totalElapsedMilliseconds;
 	}