Registered server commands - Leave the connection of the caller alone (#10577)

Author: andreasjordan

private/functions/Disconnect-Regserver.ps1

@@ -1,8 +1,22 @@
 function Disconnect-Regserver ($Server) {
+    <#
+    .SYNOPSIS
+        Internal function. Closes the connection behind a registered server object, but only if we opened it.
+
+    .DESCRIPTION
+        Takes any object of the registered server tree - a store, a group or a registered server - and walks up
+        to the RegisteredServersStore, which is the object that holds the connection.
+
+        The connection is only closed when Get-DbaRegServerStore opened it itself, which it records as
+        IsNewConnection on the store. A connection that was handed in belongs to the caller, and closing it takes
+        their session, their temp tables and their database context with it. See #10572.
+    #>
     $i = 0
-    do { $server = $server.Parent }
-    until ($null -ne $server.ServerConnection -or $i++ -gt 20)
-    if ($server.ServerConnection) {
-        $server.ServerConnection.Disconnect()
+    while ($null -ne $Server -and $null -eq $Server.ServerConnection -and $i++ -le 20) {
+        $Server = $Server.Parent
     }
-}
+
+    if ($Server.ServerConnection -and $Server.IsNewConnection) {
+        $Server.ServerConnection.Disconnect()
+    }
+}

public/Add-DbaRegServerGroup.ps1

@@ -151,9 +151,7 @@ function Add-DbaRegServerGroup {
                     $newgroup.Alter()
 
                     Get-DbaRegServerGroup -SqlInstance $currentInstance -Group (Get-RegServerGroupReverseParse -object $newgroup)
-                    if ($currentInstance.ConnectionContext) {
-                        $currentInstance.ConnectionContext.Disconnect()
-                    }
+                    Disconnect-RegServer -Server $newgroup
                 } catch {
                     Stop-Function -Message "Failed to add $reggroup" -ErrorRecord $_ -Continue
                 }

public/Get-DbaRegServer.ps1

@@ -219,7 +219,7 @@ function Get-DbaRegServer {
                 }
             } else {
                 $servers += ($serverstore.DatabaseEngineServerGroup.GetDescendantRegisteredServers())
-                $serverstore.ServerConnection.Disconnect()
+                Disconnect-RegServer -Server $serverstore
             }
 
             # save the $serverstore for later usage

public/Get-DbaRegServerGroup.ps1

@@ -177,9 +177,7 @@ function Get-DbaRegServerGroup {
                     $groups = $serverstore.DatabaseEngineServerGroup.GetDescendantRegisteredServers().Parent | Where-Object Id -In $Id
                 }
             }
-            if ($serverstore.ServerConnection) {
-                $serverstore.ServerConnection.Disconnect()
-            }
+            Disconnect-RegServer -Server $serverstore
 
             foreach ($groupobject in $groups) {
                 Add-Member -Force -InputObject $groupobject -MemberType NoteProperty -Name ComputerName -Value $serverstore.ComputerName

public/Get-DbaRegServerStore.ps1

@@ -48,7 +48,9 @@ function Get-DbaRegServerStore {
         - RegisteredServers: Collection of registered servers at the root level
 
         Properties excluded from default display (internal/technical properties):
-        - ServerConnection, DomainInstanceName, DomainName, Urn, Properties, Metadata, Parent, ConnectionContext, PropertyMetadataChanged, PropertyChanged, ParentServer
+        - ServerConnection, DomainInstanceName, DomainName, Urn, Properties, Metadata, Parent, ConnectionContext, PropertyMetadataChanged, PropertyChanged, ParentServer, IsNewConnection
+
+        IsNewConnection records whether the connection behind the store was opened here or handed in by the caller. The registered server commands use it to decide whether they may close the connection when they are done.
 
         All SMO RegisteredServersStore properties are accessible using Select-Object *, including the excluded properties if needed for advanced operations.
 
@@ -72,8 +74,17 @@ function Get-DbaRegServerStore {
     )
     process {
         foreach ($instance in $SqlInstance) {
+            # Connect-DbaInstance tells us whether it opened a connection for us. The whole registered server
+            # family closes the connection through Disconnect-RegServer, which reads the answer back off the
+            # store, so that only a connection this module opened is ever closed. See #10572.
+            $isNewConnection = $false
+            $splatConnect = @{
+                SqlInstance              = $instance
+                SqlCredential            = $SqlCredential
+                IsNewConnectionReference = [ref]$isNewConnection
+            }
             try {
-                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential
+                $server = Connect-DbaInstance @splatConnect
             } catch {
                 Stop-Function -Message "Failure" -Category ConnectionError -ErrorRecord $_ -Target $instance -Continue
             }
@@ -88,7 +99,8 @@ function Get-DbaRegServerStore {
             Add-Member -Force -InputObject $store -MemberType NoteProperty -Name InstanceName -value $server.ServiceName
             Add-Member -Force -InputObject $store -MemberType NoteProperty -Name SqlInstance -value $server.DomainInstanceName
             Add-Member -Force -InputObject $store -MemberType NoteProperty -Name ParentServer -value $server
-            Select-DefaultView -InputObject $store -ExcludeProperty ServerConnection, DomainInstanceName, DomainName, Urn, Properties, Metadata, Parent, ConnectionContext, PropertyMetadataChanged, PropertyChanged, ParentServer
+            Add-Member -Force -InputObject $store -MemberType NoteProperty -Name IsNewConnection -value $isNewConnection
+            Select-DefaultView -InputObject $store -ExcludeProperty ServerConnection, DomainInstanceName, DomainName, Urn, Properties, Metadata, Parent, ConnectionContext, PropertyMetadataChanged, PropertyChanged, ParentServer, IsNewConnection
         }
 
         # Magic courtesy of Mathias Jessen and David Shifflet

public/Move-DbaRegServer.ps1

@@ -126,7 +126,7 @@ function Move-DbaRegServer {
                 try {
                     $null = $parentserver.ServerConnection.ExecuteNonQuery($regserver.ScriptMove($movetogroup).GetScript())
                     Get-DbaRegServer -SqlInstance $server -Name $regserver.Name -ServerName $regserver.ServerName
-                    $parentserver.ServerConnection.Disconnect()
+                    Disconnect-RegServer -Server $parentserver
                 } catch {
                     Stop-Function -Message "Failed to move $($regserver.Name) to $Group on $($regserver.SqlInstance)" -ErrorRecord $_ -Continue
                 }

public/Move-DbaRegServerGroup.ps1

@@ -130,7 +130,7 @@ function Move-DbaRegServerGroup {
                     Write-Message -Level Verbose -Message "Executing $($regservergroup.ScriptMove($groupobject).GetScript())"
                     $null = $parentserver.ServerConnection.ExecuteNonQuery($regservergroup.ScriptMove($groupobject).GetScript())
                     Get-DbaRegServerGroup -SqlInstance $server -Group $newname
-                    $parentserver.ServerConnection.Disconnect()
+                    Disconnect-RegServer -Server $parentserver
                 } catch {
                     Stop-Function -Message "Failed to move $($regserver.Name) to $NewGroup on $($regserver.SqlInstance)" -ErrorRecord $_ -Continue
                 }

public/Remove-DbaRegServerGroup.ps1

@@ -115,7 +115,7 @@ function Remove-DbaRegServerGroup {
                 # try to avoid 'Collection was modified after the enumerator was instantiated' issue
                 if ($regservergroup.ID) {
                     $null = $parentserver.ServerConnection.ExecuteNonQuery($regservergroup.ScriptDrop().GetScript())
-                    $parentserver.ServerConnection.Disconnect()
+                    Disconnect-RegServer -Server $parentserver
                 } else {
                     $regservergroup.Drop()
                 }

tests/Add-DbaRegServerGroup.Tests.ps1

@@ -108,4 +108,39 @@ Describe $CommandName -Tag IntegrationTests {
             $results.SqlInstance | Should -Not -BeNullOrEmpty
         }
     }
+
+    Context "The connection of the caller is left alone (#10572)" {
+        BeforeAll {
+            # We want to run all commands in the BeforeAll block with EnableException to ensure that the test fails if the setup fails.
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            # Only a non-pooled connection can show this. SMO silently reopens a pooled connection, so the test
+            # would pass even with the disconnect this is about.
+            $callerServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -NonPooledConnection
+            $null = $callerServer.ConnectionContext.ExecuteNonQuery("CREATE TABLE #dbatoolsci_marker (id INT)")
+
+            $callerGroupName = "dbatoolsci-caller-group"
+            $callerResult = Add-DbaRegServerGroup -SqlInstance $callerServer -Name $callerGroupName
+
+            # We want to run all commands outside of the BeforeAll block without EnableException to be able to test for specific warnings.
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        AfterAll {
+            # We want to run all commands in the AfterAll block with EnableException to ensure that the test fails if the cleanup fails.
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $null = $callerServer | Disconnect-DbaInstance
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "still adds the group" {
+            $callerResult.Name | Should -Be $callerGroupName
+        }
+
+        It "leaves the connection open, so the session survives" {
+            { $callerServer.ConnectionContext.ExecuteScalar("SELECT COUNT(*) FROM #dbatoolsci_marker") } | Should -Not -Throw
+        }
+    }
 }

tests/Get-DbaRegServer.Tests.ps1

@@ -197,4 +197,44 @@ Describe $CommandName -Tag IntegrationTests {
 
         # Property Comparisons will come later when we have the commands
     }
+
+    Context "The connection of the caller is left alone (#10572)" {
+        BeforeAll {
+            # We want to run all commands in the BeforeAll block with EnableException to ensure that the test fails if the setup fails.
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            # This context brings its own registered server, because the one of the context above is cleaned up
+            # by the time this runs.
+            $callerRegSrvName = "dbatoolsci-caller-server"
+            $null = Add-DbaRegServer -SqlInstance $TestConfig.InstanceSingle -ServerName $callerRegSrvName -Name $callerRegSrvName
+
+            # Only a non-pooled connection can show this. SMO silently reopens a pooled connection, so the test
+            # would pass even with the disconnect this is about.
+            $callerServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -NonPooledConnection
+            $null = $callerServer.ConnectionContext.ExecuteNonQuery("CREATE TABLE #dbatoolsci_marker (id INT)")
+
+            $callerResult = Get-DbaRegServer -SqlInstance $callerServer
+
+            # We want to run all commands outside of the BeforeAll block without EnableException to be able to test for specific warnings.
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        AfterAll {
+            # We want to run all commands in the AfterAll block with EnableException to ensure that the test fails if the cleanup fails.
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $null = $callerServer | Disconnect-DbaInstance
+            Get-DbaRegServer -SqlInstance $TestConfig.InstanceSingle -Name $callerRegSrvName | Remove-DbaRegServer -ErrorAction SilentlyContinue
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "still returns the registered servers" {
+            $callerResult.Name | Should -Contain $callerRegSrvName
+        }
+
+        It "leaves the connection open, so the session survives" {
+            { $callerServer.ConnectionContext.ExecuteScalar("SELECT COUNT(*) FROM #dbatoolsci_marker") } | Should -Not -Throw
+        }
+    }
 }