Database context - Add Restore-DatabaseContext and fix eighteen commands

[andreasjordan]

Type of Change

• Bug fix (non-breaking change, fixes Invoke-DbaDbUpgrade - Leaves the connection pointing at the upgraded database #10556, closes Database-scoped SMO calls silently change the current database of the shared connection #10555)
• Pester test is included

This completes bucket B of #10555 and closes that issue: with it, everything the issue documents as fixable by the wrapper mechanism is merged or in here - bucket A in #10579, the server level statements in #10580, the consumer side check in #10564. The two remaining surfaces, SMO Create()/Drop() (bucket C) and collection enumeration (bucket D), continue in #10604, which starts with the scope decision to make before any batch is written.

It replaces #10602, which fixed Invoke-DbaDbUpgrade on its own. That fix is in here, using the helper below instead of its own copy of the same block, so the two cannot be merged separately in the wrong order.

Purpose

Eighteen commands handed the connection back pointing at the last database they touched. Invoke-DbaDbUpgrade is the command from the provisioning script that started this whole workstream (#10556); most of the others only read, which makes the leak worse rather than better - reading is supposed to be harmless.

Measured on SQL Server 2019 with -NonPooledConnection, before the fix. The row counts are there to show the commands were doing real work while they leaked:

Find-DbaTrigger           LEAKED   rows=1     Get-DbaPermission        LEAKED   rows=235
Find-DbaView              LEAKED   rows=1     Get-DbaLastGoodCheckDb   LEAKED   rows=1
Find-DbaStoredProcedure   LEAKED   rows=1     Get-DbaDbSpace           LEAKED   rows=2
Find-DbaObject            LEAKED   rows=4     Invoke-DbaDbUpgrade      LEAKED
Find-DbaDbDisabledIndex   LEAKED   rows=1
Get-DbaQueryExecutionTime LEAKED   rows=28

The second batch: Remove-DbaDbOrphanUser, Repair-DbaDbOrphanUser, ConvertTo-DbaXESession, Find-DbaOrphanedFile, Invoke-DbaDbDecryptObject and Copy-DbaDatabase.

The third batch closes bucket B: Copy-DbaLogin and Sync-DbaLoginPermission, together with the private Update-SqlPermission that does the permission work for both. These three are one unit on purpose - measured earlier on this issue, routing the two ExecuteNonQuery calls of Update-SqlPermission through the fixed wrapper alone does not stop the callers from leaking, because the leak is also the enumeration of database level collections ($destDb.Users[...], $destDb.Roles, $sourceDb.Roles) and, in Copy-DbaLogin, SMO's own Login.Create() and Login.Drop(), which move the connection to master. Both commands leak on two connections, the source and every destination, and both are put back.

With this, every direct SMO ExecuteNonQuery/ExecuteWithResults on a Database object in the module is either fixed or deliberately left alone: Get-EncryptedObjectImageValue keeps its call because rerouting it broke the off-row decryption path, and its only caller Invoke-DbaDbDecryptObject restores the context instead.

Approach

Every one of these commands has two causes, and fixing either alone leaves the leak in place.

1. The call. Database.ExecuteWithResults and Database.ExecuteNonQuery are SMO's own methods and issue a USE on the execution manager of the database - which is the connection context of the parent server and belongs to the caller. They go through our Query and Invoke script methods now, which restore the previous database in a finally since #10579. Statements that never needed a database context run on the server connection instead (the sys.master_files reads in Copy-DbaDatabase).

2. Enumerating a database level collection does exactly the same, and no script method can cover it, because it is a property getter and not a call we make:

db.Views                      leaked        db.Size (a plain property)   ok
db.Tables                     leaked        db.Refresh()                 ok
db.StoredProcedures           leaked
db.Users                      leaked
db.Schemas                    leaked
db.Roles                      leaked
db.FileGroups                 leaked
db.Tables[0].Columns          leaked

So each command reads ConnectionContext.CurrentDatabase before it starts any work - reading it later records a database an earlier statement already leaked into, and the restore then does nothing - and puts it back in a finally.

The new private function

That block existed by hand in Set-DbaTempDbConfig (#10580) and in the four script methods (#10579), and this change would have added many more copies. private/functions/Restore-DatabaseContext.ps1 is that block once: the case sensitive comparison, the escaping of a closing bracket in the name, and the rule that a failing restore warns rather than throwing - the database of the caller can be dropped, renamed, taken offline or made unreachable by the very statement that succeeded, and housekeeping must never replace the outcome of the command. Interrupting warning preferences (Stop, Inquire) are neutralised for that one warning, so cleanup can never terminate a command that succeeded - found in review, with its own regression test.

The script methods in xml/dbatools.Types.ps1xml keep their own copy on purpose; they are not commands and do not resolve module private functions. Set-DbaTempDbConfig was left alone, so this pull request stays about the commands it fixes.

One behaviour trap worth flagging for review

Where ExecuteWithResults was replaced, the .Tables.Rows and .Tables[0] readers had to go with it. Database.Query ends in .Tables[0], and PowerShell unrolls a DataTable on output, so the wrapper already hands back the rows:

ExecuteWithResults().Tables.Rows : count=3 firstName=master
Query()                          : count=3 firstName=master
Query().Rows                     : count=3 firstName=          <- one empty object per row

Keeping the .Rows compiles, runs, returns the right number of objects and empties every column. The existing tests of Get-DbaDbSpace and Get-DbaPermission caught it. The reverse trap exists too: ConvertTo-DbaXESession needs Query($sql, $true), because the conversion procedure spreads its script over several result sets and the one-argument form returns only the first.

Copy-DbaDatabase and Find-DbaOrphanedFile also carried $fttable = $null = ..., a double assignment that always left $fttable empty, so no full text catalog was ever copied or reported. Both sit on the SQL Server 2005 path, which no instance in the lab or in CI can reach, so they get the fix and this description but no test.

Tests

Every command got the same regression Context: its own database and objects, a non-pooled caller connection (a pooled one reconnects at its default database and hides the leak), the command run against it, then DB_NAME() compared with what it was before. Each also asserts the command still returned what it was asked for, because a command that does nothing cannot leak and would pass for the wrong reason. Get-DbaQueryExecutionTime had no integration tests at all and has some now.

Invoke-DbaDbUpgrade additionally has a -NoRefreshView -Force case, which is the one the restore around the view enumeration cannot cover, so the statements have to put the database back themselves; and a case where the caller sits in tempdb, which fails for any fix that restores "to master" instead of restoring what the caller had.

Copy-DbaLogin and Sync-DbaLoginPermission assert both connections, source and destination, each probed on its own non-pooled connection passed in as -Source and -Destination.

Restore-DatabaseContext has its own test file: the collection leak, the bracket escaping against a database really named dbatoolsci_br]acket_*, doing nothing when nothing was remembered, warning instead of throwing when the database cannot be reached - including with -WarningAction Stop - and running on the session of the caller rather than a copy.

All green in the lab, and each verified to have teeth by reverting the fix:

test file	result
Restore-DatabaseContext	6 passed
Invoke-DbaDbUpgrade	9 passed
Find-DbaTrigger	12 passed
Find-DbaView	8 passed
Find-DbaStoredProcedure	7 passed
Find-DbaObject	16 passed
Find-DbaDbDisabledIndex	6 passed
Get-DbaQueryExecutionTime	3 passed
Get-DbaPermission	10 passed
Get-DbaLastGoodCheckDb	9 passed
Get-DbaDbSpace	16 passed
Remove-DbaDbOrphanUser	11 passed
Repair-DbaDbOrphanUser	7 passed
ConvertTo-DbaXESession	4 passed
Find-DbaOrphanedFile	8 passed
Invoke-DbaDbDecryptObject	48 passed
Copy-DbaDatabase	28 passed
Copy-DbaLogin	20 passed (2 pre-existing skips)
Sync-DbaLoginPermission	7 passed

Removing the single Restore-DatabaseContext call from Find-DbaView fails its context test and nothing else. Reverting Invoke to ExecuteNonQuery in Invoke-DbaDbUpgrade fails only its -NoRefreshView case. Dropping the bracket escaping from the helper fails only the bracket test. Removing the restores from Copy-DbaLogin and Sync-DbaLoginPermission fails their context tests and nothing else.

Commands to test

$server = Connect-DbaInstance -SqlInstance $instance -NonPooledConnection
$null = Get-DbaDbSpace -SqlInstance $server -Database <any user database>
$server.ConnectionContext.ExecuteScalar("SELECT DB_NAME()")   # master, was that database before

---

This text was created by Claude and reviewed by Andreas Jordan.

[potatoqualitee]

Review verdict

Request changes. Do not merge as-is. I found one correctness issue in the new helper. After that is fixed, I would approve the PR.

Blocking: restore failure can still terminate the original command

File: private/functions/Restore-DatabaseContext.ps1

The helper explicitly promises that restoring the database is housekeeping and must never become the outcome of the command. However, its catch block currently does this:

Write-Message -Level Warning -Message "The database context could not be restored to [$Database]: $_"

Write-Message -Level Warning ultimately writes to PowerShell's warning stream. If the caller uses -WarningAction Stop or has $WarningPreference = "Stop", that warning becomes terminating. Inquire can also stop execution waiting for input.

That means the actual command can succeed, restoration can fail because the original database is no longer accessible, and the caller can receive a terminating failure caused only by cleanup. That contradicts the helper's stated contract.

The existing Database.Query and Database.Invoke restoration code already protects against this exact problem by neutralizing interrupting warning preferences before writing the warning.

The current helper test only exercises WarningAction = "SilentlyContinue", so it does not catch the terminating case.

Suggested implementation

Using a splat:

catch {
    if ($WarningPreference -notin "SilentlyContinue", "Continue") {
        $WarningPreference = "Continue"
    }

    $splatMessage = @{
        Level   = "Warning"
        Message = "The database context could not be restored to [$Database]: $_"
    }

    Write-Message @splatMessage
}

This mirrors the behavior already established by the Query/Invoke restoration code: preserve silent behavior when requested, but prevent Stop, Inquire, etc. from allowing housekeeping to replace the real command outcome.

Regression test

I would add a test specifically for WarningAction Stop:

It "does not throw when warning action is Stop" {
    & $moveTheConnection

    $missingDatabase = "dbatoolsci_does_not_exist_$(Get-Random)"

    $splatRestore = @{
        Server        = $callerServer
        Database      = $missingDatabase
        WarningAction = "Stop"
    }

    {
        Restore-DatabaseContext @splatRestore
    } | Should -Not -Throw

    $callerServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") |
        Should -Be $contextDbName
}

Copyable GitHub review comment

Restore-DatabaseContext promises that cleanup can never replace the outcome of the calling command, but Write-Message -Level Warning still honors the caller's $WarningPreference. With -WarningAction Stop, a failed restore becomes terminating; with Inquire, it can block for input.

Please normalize the warning preference here as the existing Query/Invoke restoration blocks do, and add a regression test using -WarningAction Stop. The current test uses SilentlyContinue, so it cannot catch this case.

Everything else

The overall architecture looks correct.

The original database is captured before database-scoped work begins, and restoration is correctly placed in finally, including paths containing continue or exceptions. Database names containing ] are escaped correctly, and restoration uses the caller's existing connection rather than a copied connection, preserving the SPID and session state.

The command-level regression tests are also well constructed. They use non-pooled connections so reconnect behavior cannot hide the leak, they ensure the commands actually perform work, and then verify the caller's DB_NAME() afterward.

I also checked the ExecuteWithResults → Query conversions. Removing .Tables.Rows/.Tables[0] is correct because the dbatools Database.Query script method already returns the first DataTable, which PowerShell enumerates as rows.

Invoke-DbaDbUpgrade looks correct as well: the maintenance statements now use the fixed Invoke wrapper, while the $db.Views enumeration is separately protected by the helper.

I found no other merge-blocking issue. Fix the warning-preference case and add the regression test, and I would approve #10603.

[potatoqualitee]

im going to fix this then merge on pass 👍🏼

[andreasjordan]

Thank you for the fix - measured here before it landed, and it reproduces exactly as you described: -WarningAction Stop throws The running command stopped because the preference variable "WarningPreference" or common parameter is set to Stop, and -WarningAction Inquire throws PowerShell is in NonInteractive mode on a runner. One case behaves differently, and only makes your fix more worth having rather than less: a $WarningPreference = "Stop" set in the caller's own scope did not reach the helper and returned normally, so the common parameter is the path that bites. My test only ever passed SilentlyContinue, which is why it never saw any of this.

One last commit for today follows: six more commands of bucket B, all green in the lab against SQL03\SQL2019 and SQL03\SQL2025.

command	tests
Remove-DbaDbOrphanUser	11 passed
Repair-DbaDbOrphanUser	7 passed
ConvertTo-DbaXESession	4 passed
Find-DbaOrphanedFile	8 passed
Invoke-DbaDbDecryptObject	48 passed
Copy-DbaDatabase	28 passed

If you would rather merge what is here now, say so and the commit goes into a follow-up pull request instead - it is one commit and easy to move.

Three things in it are worth a second look, because measurement contradicted the obvious change:

• ConvertTo-DbaXESession needed Query($sql, $true), not Query($sql). The conversion procedure returns the script spread over several result sets, and Query without the second argument returns only the first - which produced a script that failed with Incorrect syntax near ')'. Caught by its own existing test.
• Get-EncryptedObjectImageValue was left alone on purpose. Swapping its ExecuteWithResults(...).Tables for Query($commandText, $true) broke four off row decryption tests, so the leak is fixed in Invoke-DbaDbDecryptObject itself instead, where the connection is anyway moved by enumerating StoredProcedures, UserDefinedFunctions and Views.
• Update-SqlPermission was left alone as well. Routing its two ExecuteNonQuery calls through the wrapper does not stop Sync-DbaLoginPermission leaking - measured, still LEAKED afterwards - because the leak is collection enumeration in the calling commands. It belongs with Copy-DbaLogin and Sync-DbaLoginPermission rather than half done here.

Copy-DbaDatabase also carries the $fttable = $null = ... fix, the same double assignment as in Find-DbaOrphanedFile. Both sit on the SQL Server 2005 path, which no instance in the lab or in CI can reach, so they get the fix and this description but no test.

---

This text was created by Claude and reviewed by Andreas Jordan.

[andreasjordan]

I have to stop here for today.

[potatoqualitee]

sounds good, we ready to publish after this merge?

[andreasjordan]

No, this is only (part of?) bucket B - and there is Bucket C and maybe more. I would vote for fixing avery part of #10555 and be really sure this works.

[andreasjordan]

The last of bucket B is in: Copy-DbaLogin and Sync-DbaLoginPermission

One more commit, and bucket B of #10555 is closed. These two commands and the private Update-SqlPermission that does the permission work for both are one unit on purpose - measured on the issue, routing the two ExecuteNonQuery calls of Update-SqlPermission through the fixed wrapper alone does not stop the callers from leaking, because the leak is also the enumeration of database level collections ($destDb.Users, $destDb.Roles, $sourceDb.Roles) and, in Copy-DbaLogin, SMO's own Login.Create() and Login.Drop(), which move the connection to master.

Two things are new compared to the commands already in here:

• Both commands leak on two connections, the source and every destination. Both are captured before any work and both are restored in a finally - each destination after its iteration, the source after all destinations. The regression tests probe both, each on its own non-pooled connection passed in as -Source and -Destination, with a database user mapping so the permission sync really walks the database level collections.
• Get-EncryptedObjectImageValue stays as it is (see the earlier comment): with this commit, every direct SMO ExecuteNonQuery/ExecuteWithResults on a Database object in the module is either fixed or deliberately left alone.

All green in the lab against SQL03\SQL2022 and SQL04\SQL2025 (COPY set) and SQL03\SQL2025 / SQL03\SQL2022 (MULTI set), no warnings, environment left clean:

test file	result
Copy-DbaLogin	20 passed (2 pre-existing skips)
Sync-DbaLoginPermission	7 passed

Verified to have teeth by stashing the fix and re-running: exactly the four context assertions fail, each reporting the leaked database (Expected: 'master' But was: 'dbatoolsci_ctx_copylogin_...'), and nothing else.

The title and body of this pull request now describe all eighteen commands, so the squash commit tells the whole story.

@potatoqualitee this is ready for review from my side - buckets C and D of #10555 stay open, and I would discuss how far we take those before the next batch starts.

---

This text was created by Claude and reviewed by Andreas Jordan.

[andreasjordan]

@potatoqualitee I think this can be merged now. All open topics are moved to #10604.