Connect-DbaInstance / Invoke-DbaQuery - Do not close connections of the caller and do not reuse them for the wrong database

public/Connect-DbaInstance.ps1

@@ -157,6 +157,10 @@ function Connect-DbaInstance {
         Creates a dedicated administrator connection (DAC) for emergency access to SQL Server.
         Use this when SQL Server is unresponsive to regular connections, allowing you to diagnose and resolve critical issues. Remember to manually disconnect the connection when finished.
 
+    .PARAMETER IsNewConnectionReference
+        Reports whether a new connection was opened, by writing $true or $false into the referenced variable: pass a variable that already exists as [ref]$variable.
+        Use this when your command has to clean up after itself: only close a connection when this is $true, because a connection that was passed in belongs to the caller and closing it takes their session, their temp tables and their database context with it. When several instances are connected in one call, the value reflects the last connection that was returned.
+
     .PARAMETER DisableException
         Changes exception handling from throwing errors to displaying warnings.
         Use this in interactive sessions where you want graceful error handling instead of script-stopping exceptions, which is the default behavior for this command.
@@ -414,6 +418,7 @@ function Connect-DbaInstance {
         [ValidateSet('ActiveDirectoryIntegrated', 'ActiveDirectoryInteractive', 'ActiveDirectoryPassword', 'ActiveDirectoryServicePrincipal', 'ActiveDirectoryManagedIdentity', 'ActiveDirectoryDeviceCodeFlow')]
         [string]$AuthenticationType,
         [switch]$DedicatedAdminConnection,
+        [ref]$IsNewConnectionReference,
         [switch]$DisableException
     )
     begin {
@@ -1261,6 +1266,9 @@ function Connect-DbaInstance {
                     $null = Add-ConnectionHashValue -Key $server.ConnectionContext.ConnectionString -Value $server.ConnectionContext.SqlConnectionObject
                 }
                 Write-Message -Level Debug -Message "We return only SqlConnection in server.ConnectionContext.SqlConnectionObject"
+                if ($IsNewConnectionReference) {
+                    $IsNewConnectionReference.Value = $isNewConnection
+                }
                 $server.ConnectionContext.SqlConnectionObject
                 continue
             }
@@ -1313,6 +1321,9 @@ function Connect-DbaInstance {
             }
 
             Write-Message -Level Debug -Message "We return the server object"
+            if ($IsNewConnectionReference) {
+                $IsNewConnectionReference.Value = $isNewConnection
+            }
             $server
 
             if ($isNewConnection -and -not $DedicatedAdminConnection) {

public/Invoke-DbaQuery.ps1

@@ -496,23 +496,34 @@ function Invoke-DbaQuery {
                 # Again, here we are, but we cannot use the connection when an information is lost, which is that we are:
                 #   - Integrated Security=True
                 #   - using ConnectAsUser, ConnectAsUserName, ConnectAsUserPassword to "log on as a different user"
+                #
+                # The database is tested with ConnectionContext.CurrentDatabase, the database the connection is on
+                # right now, and not with ConnectionContext.DatabaseName, which is only the database the connection
+                # was opened with. The two differ as soon as anything runs a USE, and then reusing the connection
+                # runs the query in the wrong database. Connect-DbaInstance tests CurrentDatabase as well, so both
+                # commands now ask the same question.
 
                 $startedWithAnOpenConnection = # we want to bypass Connect-DbaInstance if
                 ($instance.InputObject.GetType().Name -eq "Server") -and # we have Server SMO object and
                 (-not $ReadOnly) -and # no readonly intent is requested and
-                (-not $Database -or $instance.InputObject.ConnectionContext.DatabaseName -eq $Database) -and # the database is not set or the currently connected and
+                (-not $Database -or $instance.InputObject.ConnectionContext.CurrentDatabase -eq $Database) -and # the database is not set or the connection is on it and
                 (-not $AppendConnectionString) -and # we don't use AppendConnectionString and
                 ($instance.InputObject.ConnectionContext.ConnectAsUserName -eq '') # we don't use a DIFFERENT operating system user than the current logged in
+                # Connect-DbaInstance tells us whether it opened a connection for us. We must only close what we opened
+                # ourselves, because closing a connection of the caller takes their session with it. As we might not
+                # call Connect-DbaInstance at all, we have to set the default here.
+                $isNewConnection = $false
                 if ($startedWithAnOpenConnection) {
                     Write-Message -Level Debug -Message "Current connection will be reused"
                     $server = $instance.InputObject
                 } else {
                     $connDbaInstanceParams = @{
-                        SqlInstance         = $instance
-                        SqlCredential       = $SqlCredential
-                        Database            = $Database
-                        NonPooledConnection = $true           # see #8491 for details, also #7725 is still relevant
-                        Verbose             = $false
+                        SqlInstance              = $instance
+                        SqlCredential            = $SqlCredential
+                        Database                 = $Database
+                        NonPooledConnection      = $true           # see #8491 for details, also #7725 is still relevant
+                        IsNewConnectionReference = [ref]$isNewConnection
+                        Verbose                  = $false
                     }
                     if ($ReadOnly) {
                         $connDbaInstanceParams.ApplicationIntent = "ReadOnly"
@@ -541,9 +552,11 @@ function Invoke-DbaQuery {
             } catch {
                 Stop-Function -Message "[$instance] Failed during execution" -ErrorRecord $_ -Target $instance -Continue
             }
-            # if the given connection started out open, don't close it.
-            if ($connDbaInstanceParams.NonPooledConnection -and -not $startedWithAnOpenConnection) {
-                # Close non-pooled connection as this is not done automatically. If it is a reused Server SMO, connection will be opened again automatically on next request.
+            # Only close the connection if Connect-DbaInstance opened a new one for us. Connect-DbaInstance returns the
+            # object that was passed in whenever nothing about it has to change, so testing our own parameters is not
+            # enough - see #10554.
+            if ($isNewConnection) {
+                # Close non-pooled connection as this is not done automatically.
                 $null = $server | Disconnect-DbaInstance -Verbose:$false
             }
         }

tests/Connect-DbaInstance.Tests.ps1

@@ -48,6 +48,7 @@ Describe $CommandName -Tag UnitTests {
                 "AccessToken",
                 "AuthenticationType",
                 "DedicatedAdminConnection",
+                "IsNewConnectionReference",
                 "DisableException"
             )
             Compare-Object -ReferenceObject $expectedParameters -DifferenceObject $hasParameters | Should -BeNullOrEmpty
@@ -418,6 +419,46 @@ Describe $CommandName -Tag IntegrationTests {
         }
     }
 
+    Context "IsNewConnectionVariable tells the caller whether a connection was opened" {
+        BeforeAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            # We predefine the variables so that a test fails if Connect-DbaInstance does not set them at all.
+            $newFromString = $null
+            $newFromServer = $null
+            $newFromCopy = $null
+
+            $serverFromString = Connect-DbaInstance -SqlInstance $TestConfig.InstanceMulti1 -NonPooledConnection -IsNewConnectionReference ([ref]$newFromString)
+            $serverFromServer = Connect-DbaInstance -SqlInstance $serverFromString -IsNewConnectionReference ([ref]$newFromServer)
+            # Asking for a different database forces Connect-DbaInstance to copy the connection context, so this is a new connection.
+            $serverFromCopy = Connect-DbaInstance -SqlInstance $serverFromString -Database tempdb -IsNewConnectionReference ([ref]$newFromCopy)
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        AfterAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+            $null = $serverFromString, $serverFromCopy | Disconnect-DbaInstance
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "is true when the connection is opened from a string" {
+            $newFromString | Should -BeTrue
+        }
+
+        It "is false when the server object is passed back in" {
+            $newFromServer | Should -BeFalse
+        }
+
+        It "returns the object that was passed in when nothing has to change" {
+            [object]::ReferenceEquals($serverFromServer, $serverFromString) | Should -BeTrue
+        }
+
+        It "is true when the connection context has to be copied" {
+            $newFromCopy | Should -BeTrue
+        }
+    }
+
     Context "connection is properly made using a connection string" {
         BeforeAll {
             $server = Connect-DbaInstance -SqlInstance "Data Source=$($TestConfig.InstanceMulti1);Initial Catalog=tempdb;Integrated Security=True"

tests/Invoke-DbaQuery.Tests.ps1

@@ -401,4 +401,88 @@ CREATE INDEX IX_Filtered ON dbo.$tableName(Name) WHERE IsDeleted = 0;
         $results = Invoke-DbaQuery -SqlInstance $TestConfig.InstanceMulti1 -Query "select cast(null as hierarchyid)"
         $results.Column1 | Should -Be "NULL"
     }
+
+    Context "The connection is only reused when it is on the requested database (#10554)" {
+        BeforeAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $movedServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceMulti1 -Database tempdb -NonPooledConnection
+            $movedOwnSpid = $movedServer.ConnectionContext.ExecuteScalar("SELECT @@SPID")
+
+            # As long as the connection is still on tempdb, it is reused.
+            $reusedSpid = Invoke-DbaQuery -SqlInstance $movedServer -Database tempdb -Query "SELECT @@SPID AS spid" -As SingleValue
+
+            # A USE moves the connection to another database. ConnectionContext.DatabaseName still says tempdb,
+            # only ConnectionContext.CurrentDatabase knows that the connection is on master now.
+            $null = $movedServer.ConnectionContext.ExecuteNonQuery("USE [master]")
+            $movedDatabase = Invoke-DbaQuery -SqlInstance $movedServer -Database tempdb -Query "SELECT DB_NAME() AS dbname" -As SingleValue
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        AfterAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $null = $movedServer | Disconnect-DbaInstance
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "reuses the connection while it is on the requested database" {
+            $reusedSpid | Should -Be $movedOwnSpid
+        }
+
+        It "runs in the requested database after the connection was moved away from it" {
+            $movedDatabase | Should -Be "tempdb"
+        }
+    }
+
+    Context "Connections that were passed in are not closed (#10554)" {
+        BeforeAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $callerServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceMulti1 -NonPooledConnection
+            $null = $callerServer.ConnectionContext.ExecuteNonQuery("CREATE TABLE #dbatoolsci_marker (id INT)")
+
+            # Naming the database the connection is already on is enough to take the path where the connection of the caller used to be closed.
+            $null = Invoke-DbaQuery -SqlInstance $callerServer -Database master -Query "SELECT 1"
+
+            # Every call with a string opens and closes a connection of its own, so we count the sessions to see that they are still closed.
+            $counterServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceMulti1
+            $splatCountSessions = @{
+                SqlInstance  = $counterServer
+                Query        = "SELECT COUNT(*) FROM sys.dm_exec_sessions WHERE program_name = @clientName"
+                SqlParameter = @{ clientName = Get-DbatoolsConfigValue -FullName sql.connection.clientname }
+                As           = "SingleValue"
+            }
+            $sessionsBefore = Invoke-DbaQuery @splatCountSessions
+            foreach ($run in 1..5) {
+                $null = Invoke-DbaQuery -SqlInstance $TestConfig.InstanceMulti1 -Query "SELECT 1"
+            }
+            $sessionsAfter = Invoke-DbaQuery @splatCountSessions
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        AfterAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $null = $callerServer, $counterServer | Disconnect-DbaInstance
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "leaves the connection of the caller open" {
+            $callerServer.ConnectionContext.IsOpen | Should -BeTrue
+        }
+
+        It "leaves the session of the caller intact, so the temp table is still there" {
+            { $callerServer.ConnectionContext.ExecuteScalar("SELECT COUNT(*) FROM #dbatoolsci_marker") } | Should -Not -Throw
+        }
+
+        It "still closes the connections it opens itself (#6210)" {
+            # We allow for a little noise, because the tab expansion of dbatools connects in the background as well.
+            ($sessionsAfter - $sessionsBefore) | Should -BeLessThan 5
+        }
+    }
 }