Registered server commands - Leave the connection of the caller alone

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
+        }
+    }
 }

tests/Get-DbaRegServerGroup.Tests.ps1

@@ -132,4 +132,38 @@ 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
+
+            # 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-DbaRegServerGroup -SqlInstance $callerServer -Id 1
+
+            # 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 returns the group" {
+            $callerResult.Name | Should -Be "DatabaseEngineServerGroup"
+        }
+
+        It "leaves the connection open, so the session survives" {
+            { $callerServer.ConnectionContext.ExecuteScalar("SELECT COUNT(*) FROM #dbatoolsci_marker") } | Should -Not -Throw
+        }
+    }
 }

tests/Get-DbaRegServerStore.Tests.ps1

@@ -28,4 +28,37 @@ Describe $CommandName -Tag IntegrationTests {
             $results.DisplayName | Should -Be "Central Management Servers"
         }
     }
+
+    Context "The store records who owns the connection (#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 is what the whole registered server family reads to decide whether it may close the
+            # connection when it is done.
+            $callerServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -NonPooledConnection
+            $storeFromName = Get-DbaRegServerStore -SqlInstance $TestConfig.InstanceSingle
+            $storeFromServer = Get-DbaRegServerStore -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
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "reports a connection it opened itself as its own" {
+            $storeFromName.IsNewConnection | Should -BeTrue
+        }
+
+        It "reports a connection of the caller as not its own" {
+            $storeFromServer.IsNewConnection | Should -BeFalse
+        }
+    }
 }

tests/Move-DbaRegServer.Tests.ps1

@@ -89,4 +89,49 @@ Describe $CommandName -Tag IntegrationTests {
         $results = Get-DbaRegServer -SqlInstance $TestConfig.InstanceSingle -Group $testGroupHR | Move-DbaRegServer -Group $testGroupFinance
         $results.Count | Should -Be 2
     }
+
+    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
+
+            $callerRegSrvName = "dbatoolsci-caller-server"
+            $callerTargetName = "dbatoolsci-caller-target"
+            $null = Add-DbaRegServerGroup -SqlInstance $TestConfig.InstanceSingle -Name $callerTargetName
+            $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
+
+            # The marker is created after the lookup on purpose, because Get-DbaRegServer closes the connection
+            # too. This test has to measure the command under test, not its input.
+            $callerInputObject = Get-DbaRegServer -SqlInstance $callerServer -Name $callerRegSrvName
+            $null = $callerServer.ConnectionContext.ExecuteNonQuery("CREATE TABLE #dbatoolsci_marker (id INT)")
+
+            $callerResult = Move-DbaRegServer -InputObject $callerInputObject -Group $callerTargetName
+
+            # 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
+            Get-DbaRegServerGroup -SqlInstance $TestConfig.InstanceSingle -Group $callerTargetName | Remove-DbaRegServerGroup -ErrorAction SilentlyContinue
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "still moves the registered server" {
+            $callerResult.Parent.Name | Should -Be $callerTargetName
+        }
+
+        It "leaves the connection open, so the session survives" {
+            { $callerServer.ConnectionContext.ExecuteScalar("SELECT COUNT(*) FROM #dbatoolsci_marker") } | Should -Not -Throw
+        }
+    }
 }