From f0d3a0d546509dfb56587f8b6b249d15b2dc7daf Mon Sep 17 00:00:00 2001
From: Christopher Suh <chsuh@microsoft.com>
Date: Mon, 9 Oct 2023 11:18:51 -0700
Subject: [PATCH] Transfer URI Correctly (#2265)

* transfer uri regardless of queryrunner presence

* remove extra space
---
 .../Connection/ConnectionService.cs                             | 1 +
 .../QueryExecution/QueryExecutionService.cs                     | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
index f6e0217d..0a0b6609 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs
@@ -925,6 +925,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
             {
                 return false;
             }
+            info.OwnerUri = newOwnerUri;
             OwnerToConnectionMap.Remove(originalOwnerUri);
             OwnerToConnectionMap.Add(newOwnerUri, info);
             return true;
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
index e9c4bd3a..32bf7a67 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs
@@ -363,13 +363,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             {
                 string OriginalOwnerUri = changeUriParams.OriginalOwnerUri;
                 string NewOwnerUri = changeUriParams.NewOwnerUri;
+                ConnectionService.ReplaceUri(OriginalOwnerUri, NewOwnerUri);
                 // Attempt to load the query
                 Query query;
                 if (!ActiveQueries.TryRemove(OriginalOwnerUri, out query))
                 {
                     throw new Exception("Uri: " + OriginalOwnerUri + " is not associated with an active query.");
                 }
-                ConnectionService.ReplaceUri(OriginalOwnerUri, NewOwnerUri);
                 query.ConnectionOwnerURI = NewOwnerUri;
                 ActiveQueries.TryAdd(NewOwnerUri, query);
                 return Task.FromResult(true);