Get-DbaDbExtentDiff - Leave the connection of the caller alone

public/Get-DbaDbExtentDiff.ps1

@@ -105,8 +105,17 @@ function Get-DbaDbExtentDiff {
     process {
 
         foreach ($instance in $SqlInstance) {
+            # 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. See #10554.
+            $isNewConnection = $false
+            $splatConnect = @{
+                SqlInstance              = $instance
+                SqlCredential            = $SqlCredential
+                NonPooledConnection      = $true
+                IsNewConnectionReference = [ref]$isNewConnection
+            }
             try {
-                $server = Connect-DbaInstance -SqlInstance $instance -SqlCredential $SqlCredential -NonPooledConnection
+                $server = Connect-DbaInstance @splatConnect
             } catch {
                 Stop-Function -Message "Error occurred while establishing connection to $instance" -Category ConnectionError -ErrorRecord $_ -Target $instance -Continue
             }
@@ -133,14 +142,17 @@ function Get-DbaDbExtentDiff {
             #Available from 2016 SP2
             if ($server.Version -ge [version]'13.0.5026') {
                 foreach ($db in $sourcedbs) {
+                    # The database is named in the query instead of running the query in the database, because
+                    # running it in the database leaves the connection there and it belongs to the caller. See #10555.
+                    $dbNameEscaped = $db.Name.Replace("]", "]]")
                     $DBCCPageQueryDMV = "
                         SELECT
                         SUM(total_page_count) / 8 AS [ExtentsTotal],
                         SUM(modified_extent_page_count) / 8 AS [ExtentsChanged],
                         100.0 * SUM(modified_extent_page_count)/SUM(total_page_count) AS [ChangedPerc]
-                        FROM sys.dm_db_file_space_usage
+                        FROM [$dbNameEscaped].sys.dm_db_file_space_usage
                     "
-                    $DBCCPageResults = $server.Query($DBCCPageQueryDMV, $db.Name)
+                    $DBCCPageResults = $server.Query($DBCCPageQueryDMV)
                     [PSCustomObject]@{
                         ComputerName   = $server.ComputerName
                         InstanceName   = $server.ServiceName
@@ -187,8 +199,10 @@ function Get-DbaDbExtentDiff {
                 }
             }
 
-            # 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.
-            $null = $server | Disconnect-DbaInstance
+            if ($isNewConnection) {
+                # Close non-pooled connection as this is not done automatically.
+                $null = $server | Disconnect-DbaInstance
+            }
         }
     }
 }

tests/Get-DbaDbExtentDiff.Tests.ps1

@@ -70,6 +70,35 @@ Describe $CommandName -Tag IntegrationTests {
         }
     }
 
+    Context "The connection of the caller is left alone (#10554)" {
+        BeforeAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $callerServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -NonPooledConnection
+            $null = $callerServer.ConnectionContext.ExecuteNonQuery("CREATE TABLE #dbatoolsci_marker (id INT)")
+
+            $null = Get-DbaDbExtentDiff -SqlInstance $callerServer -Database $dbname
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        AfterAll {
+            $PSDefaultParameterValues["*-Dba*:EnableException"] = $true
+
+            $null = $callerServer | Disconnect-DbaInstance
+
+            $PSDefaultParameterValues.Remove("*-Dba*:EnableException")
+        }
+
+        It "leaves the connection open, so the session survives" {
+            { $callerServer.ConnectionContext.ExecuteScalar("SELECT COUNT(*) FROM #dbatoolsci_marker") } | Should -Not -Throw
+        }
+
+        It "leaves the connection in the database it was on" {
+            $callerServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "master"
+        }
+    }
+
     Context "Gets Changed Extents for Single Database" {
         BeforeAll {
             $singleDbResults = Get-DbaDbExtentDiff -SqlInstance $TestConfig.InstanceSingle -Database $dbname