Database context - Put the database back in the Query and Invoke script methods

[andreasjordan]

Type of Change

• Bug fix (non-breaking change, fixes Database-scoped SMO calls silently change the current database of the shared connection #10555 in part)
• Ran manual Pester test and has passed
• Pester test is included

Purpose

Step one of #10555: the mechanism, on its own.

The Query and Invoke script methods of Server and Database in xml/dbatools.Types.ps1xml do not run on a private connection. The execution manager of an SMO database is the connection context of the parent server, which belongs to the caller, so they issue a USE and never switch back:

call                                       result    DB_NAME() afterwards
$db.Query()                                LEAKED    dbatoolsci_wrap
$db.Invoke()                               LEAKED    dbatoolsci_wrap
$server.Query(sql) - one argument          ok        master
$server.Query(sql, db) - two arguments     LEAKED    dbatoolsci_wrap
$server.Invoke(sql, db) - two arguments    LEAKED    dbatoolsci_wrap

These four methods are reached from roughly 68 database-scoped .Query() / .Invoke() call sites plus 45 two-argument $server.Query($sql, $db) calls across 28 files, so this one file is the cheapest place in the module to fix it.

Approach

Each method remembers ConnectionContext.CurrentDatabase and puts it back in a finally, so a query that throws restores the context as well.

The Server pair needs its own copy of that. Server.Query and Server.Invoke call $this.Databases[$Database].ExecuteWithResults(...) directly and never go through the Database methods, so fixing Database.Query alone does not reach them. That is four edits, not two - the issue body assumed otherwise.

Restoring, not copying. ConnectionContext.Copy().GetDatabaseConnection($name) also works, and was the other candidate, but it is a different session:

=== copied connection context ===
    query ran in          : dbatoolsci_ctx  on SPID 64      <-- caller is on SPID 62
    copy sees temp table  : no (separate session)

=== remember and put back ===
    query ran in          : dbatoolsci_ctx  on SPID 62
    caller temp table     : still there

A copy cannot see the temp tables or SET options of the caller and opens a connection per call, which would be a silent behaviour change for any command that builds session state and then queries through the wrapper. Restoring keeps the session and costs one round trip, and only when the database actually moved.

The caller's database is restored, not master. A connection sitting in msdb is returned to msdb. Restoring to master would have passed every other test and still been wrong - and it matters, because the Agent commands move the context to msdb rather than master.

The database name is escaped for the USE, so a database containing ] in its name is handled.

Commands to test

$server = Connect-DbaInstance -SqlInstance $instance -NonPooledConnection
$null = $server.Databases["SomeDatabase"].Query("SELECT 1")
$server.ConnectionContext.ExecuteScalar("SELECT DB_NAME()")   # master, was: SomeDatabase

Tests

tests\InModule.TypeExtensions.Tests.ps1, following the existing InModule.* naming for test files that are not the test of a single command. 10 tests on InstanceSingle:

• each of the four methods leaves the database context alone
• the query still runs in the database that was asked for, and AllTables still returns every table
• the session is kept, so a temporary object created before the call is still there afterwards
• a failing query still puts the database back
• a caller connected to msdb is returned to msdb, not to master

All 10 pass. Against development 7 of them fail; the 3 that pass are the correctness assertions, which are there to catch the fix breaking something rather than to prove the bug.

Because this reaches every command that uses the wrappers, 15 further test files of wrapper-using commands were run on top: Find-DbaSimilarTable, Get-DbaCpuRingBuffer, Get-DbaDatabase, Get-DbaDbFeatureUsage, Get-DbaDbFile, Get-DbaDbSnapshot, Get-DbaDbVirtualLogFile, Get-DbaHelpIndex, Get-DbaInstanceInstallDate, Get-DbaModule, Get-DbaSchemaChangeHistory, Install-DbaWhoIsActive, Invoke-DbaDbClone, New-DbaLinkedServer, Set-DbaDbFileGrowth. 104 tests, no failures.

What this does not fix

Only the script methods. The other two sources in #10555 are untouched and still leak:

• the 35 direct $db.ExecuteNonQuery(...) / $db.ExecuteWithResults(...) call sites, which are SMO's own methods and cannot be shadowed
• SMO's own Create() and Drop() of server-level objects

Invoke-DbaDbUpgrade (#10556) is in the first of those groups. Verified against a database forced to compatibility level 100 so the upgrade really ran - it went to 150 and the connection was still left in the upgraded database.

🤖 Generated with Claude Code