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

[potatoqualitee]

gonna put this through a round of ChatGPT Pro since it's such a sensitive change

[potatoqualitee]

Glad I checked!

Findings

P1: The restore check fails on case-sensitive SQL Server instances

File: xml/dbatools.Types.ps1xml
Lines: 25, 47, 80, and 122

All four methods use this comparison:

$connectionContext.CurrentDatabase -ne $previousDatabase

PowerShell string comparison operators are case-insensitive unless the -c* form is used. SQL Server database names are instance-level identifiers and use the instance collation. On a case-sensitive instance, AppDb and appdb can be different databases. In that situation, SMO can move the connection from one to the other, but this comparison evaluates as equal and skips the restore. The original context leak therefore remains on a valid SQL Server configuration. ([Microsoft Learn]1)

At minimum, use -cne:

if ($previousDatabase -and $connectionContext.CurrentDatabase -cne $previousDatabase) {

An ordinal comparison is more explicit:

$databaseChanged = -not [string]::Equals(
    $connectionContext.CurrentDatabase,
    $previousDatabase,
    [System.StringComparison]::Ordinal
)

A server-collation-aware comparison would be exact, but an ordinal comparison is safe here. At worst it performs an unnecessary restore; it does not miss a real database change. This also needs a regression test using database names that differ only by case on a case-sensitive instance.

---

P2: A restoration failure can hide the real error, or make a successful command appear to have failed

File: xml/dbatools.Types.ps1xml
Lines: 24–28, 46–50, 79–83, and 121–125

The restoration command is unguarded inside each finally:

$null = $connectionContext.ExecuteNonQuery("USE [$escapedDatabase]")

There are two problematic outcomes:

1. The query fails, then the restore also fails. The restore exception replaces the original query exception.
2. The command succeeds but makes the previous database unavailable, for example by dropping it, taking it offline, renaming it, or revoking access. The wrapper then throws during restoration even though the requested SQL already completed.

USE requires CONNECT permission and can legitimately fail. This is especially dangerous for Invoke, because a caller may interpret the exception as “the command did not execute” and retry an operation that already succeeded. ([Microsoft Learn]2)

The wrapper should retain the original ErrorRecord as the primary failure. If the operation succeeded but restoration failed, it should throw a distinct error stating that the SQL completed but the connection context could not be restored. If both fail, preserve both errors without replacing the original.

---

P2: The temporary-table test does not prove that the wrapper uses the same session

File: tests/InModule.TypeExtensions.Tests.ps1
Lines: 81–84

The test currently:

1. Creates the temporary table through the caller connection.
2. Executes Database.Query("SELECT 1").
3. Checks for the temporary table through the caller connection again.

A copied-connection implementation would also pass this test. The temporary table remains on the original connection regardless of which session executed SELECT 1. The test therefore does not enforce the same-session behavior that motivated the implementation.

Query the temporary table through the wrapper:

$null = $callerServer.ConnectionContext.ExecuteNonQuery(
    "CREATE TABLE #dbatoolsci_marker (id INT)"
)

$result = $callerServer.Databases[$contextDbName].Query(
    "SELECT OBJECT_ID('tempdb..#dbatoolsci_marker') AS object_id"
)

$result.object_id | Should -Not -BeNullOrEmpty

Or compare @@SPID directly:

$callerSpid = $callerServer.ConnectionContext.ExecuteScalar("SELECT @@SPID")

$querySpid = $callerServer.Databases[$contextDbName].Query(
    "SELECT @@SPID AS spid"
).spid

$querySpid | Should -Be $callerSpid

Verdict

Request changes. Restoring the original context on the existing session is the correct overall approach, and the database-name escaping is correct. However, the case-insensitive comparison leaves the original bug unfixed on case-sensitive instances. The cleanup error handling also introduces ambiguous and potentially dangerous failure reporting. The temporary-table test should be corrected so it actually locks in the same-session guarantee.