Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andreasjordan wants to merge 2 commits into
base: development
Choose a base branch
from
+302 −5
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d703df9
Database context - Put the database back in the Query and Invoke scri…
andreasjordan Aug 14, 2026
cfce992
Database context - Compare case sensitively and never let the restore…
andreasjordan Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
220 changes: 220 additions & 0 deletions tests/InModule.TypeExtensions.Tests.ps1
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,220 @@
#Requires -Module @{ ModuleName="Pester"; ModuleVersion="5.0" }
param(
$ModuleName = "dbatools",
$CommandName = "InModule.TypeExtensions",
$PSDefaultParameterValues = $TestConfig.Defaults
)

Describe $CommandName -Tag IntegrationTests {
# The Query and Invoke script methods of Server and Database in xml\dbatools.Types.ps1xml run through
# the execution manager of a database, which is the connection context of the parent server and belongs
# to the caller. SMO issues a USE and never switches back, so the methods put the database back. See #10555.
BeforeDiscovery {
# Two databases whose names differ only in case can only exist on an instance with a case sensitive
# collation, and that is the only place where a case insensitive comparison in the restore can be
# caught. The collation decides it, not the version, so this is asked of the instance rather than
# assumed. On a case insensitive instance the scenario cannot be built at all and the Context skips.
$caseDiscoveryServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle
$instanceIsCaseSensitive = $caseDiscoveryServer.Collation -match "_CS_"
$null = $caseDiscoveryServer | Disconnect-DbaInstance
}

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

$contextDbName = "dbatoolsci_typeext_$(Get-Random)"
$null = New-DbaDatabase -SqlInstance $TestConfig.InstanceSingle -Name $contextDbName

# 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 = Remove-DbaDatabase -SqlInstance $TestConfig.InstanceSingle -Database $contextDbName -ErrorAction SilentlyContinue

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

Context "The database context of the caller survives the script methods (#10555)" {
BeforeAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

# Only a non-pooled connection can show this. SMO reopens a pooled connection at its default
# database, which hides the leak.
$callerServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -NonPooledConnection

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

AfterAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

$null = $callerServer | Disconnect-DbaInstance

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

It "leaves the database context alone in Database.Query" {
$null = $callerServer.Databases[$contextDbName].Query("SELECT 1")
$callerServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "master"
}

It "leaves the database context alone in Database.Invoke" {
$null = $callerServer.Databases[$contextDbName].Invoke("SELECT 1")
$callerServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "master"
}

It "leaves the database context alone in Server.Query with a database" {
$null = $callerServer.Query("SELECT 1", $contextDbName)
$callerServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "master"
}

It "leaves the database context alone in Server.Invoke with a database" {
$null = $callerServer.Invoke("SELECT 1", $contextDbName)
$callerServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "master"
}

It "still runs the query in the database that was asked for" {
$callerServer.Databases[$contextDbName].Query("SELECT DB_NAME() AS dbname").dbname | Should -Be $contextDbName
$callerServer.Query("SELECT DB_NAME() AS dbname", $contextDbName).dbname | Should -Be $contextDbName
}

It "still returns every table when AllTables is used" {
$allTables = $callerServer.Databases[$contextDbName].Query("SELECT 1 AS a; SELECT 2 AS b", $true)
$allTables.Count | Should -Be 2
}

It "runs on the session of the caller and not on a copy of the connection" {
# Checking the temporary table through the caller afterwards proves nothing: it never left the
# caller's session, so a copied connection context would pass that too. The wrapper itself has
# to see the table, and its SPID has to be the caller's. ConnectionContext.Copy() plus
# GetDatabaseConnection() was the other candidate for this fix and fails both assertions.
$callerSpid = $callerServer.ConnectionContext.ExecuteScalar("SELECT @@SPID")
$null = $callerServer.ConnectionContext.ExecuteNonQuery("CREATE TABLE #dbatoolsci_marker (id INT)")

$wrapperResult = $callerServer.Databases[$contextDbName].Query("SELECT @@SPID AS spid, (SELECT COUNT(*) FROM #dbatoolsci_marker) AS marker")

$wrapperResult.spid | Should -Be $callerSpid
$wrapperResult.marker | Should -Be 0
}

It "puts the database back even when the query fails" {
{ $callerServer.Databases[$contextDbName].Query("SELECT * FROM dbatoolsci_does_not_exist") } | Should -Throw
$callerServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "master"
}
}

Context "The database the caller was on is restored, not master (#10555)" {
BeforeAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

# A connection that starts out somewhere other than master. Restoring to master would pass the
# tests above and still be wrong here.
$msdbServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -Database msdb -NonPooledConnection

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

AfterAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

$null = $msdbServer | Disconnect-DbaInstance

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

It "leaves the connection in msdb after Database.Query" {
$null = $msdbServer.Databases[$contextDbName].Query("SELECT 1")
$msdbServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "msdb"
}

It "leaves the connection in msdb after Server.Query with a database" {
$null = $msdbServer.Query("SELECT 1", $contextDbName)
$msdbServer.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -Be "msdb"
}
}

Context "A failing restore does not become the outcome of the call (#10555)" {
BeforeAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

$offlineDbName = "dbatoolsci_typeext_offline_$(Get-Random)"
$null = New-DbaDatabase -SqlInstance $TestConfig.InstanceSingle -Name $offlineDbName
$offlineServer = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -Database $offlineDbName -NonPooledConnection

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

AfterAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

$null = $offlineServer | Disconnect-DbaInstance
$null = Set-DbaDbState -SqlInstance $TestConfig.InstanceSingle -Database $offlineDbName -Online -Force -ErrorAction SilentlyContinue
$null = Remove-DbaDatabase -SqlInstance $TestConfig.InstanceSingle -Database $offlineDbName -ErrorAction SilentlyContinue

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

It "returns normally when the statement made the previous database unreachable" {
# The wrapper moves the connection to master to run this, and the USE that would put it back
# cannot work afterwards, because by then the database it names is offline. Restoring is
# housekeeping and must not turn a statement that succeeded into an exception: a caller reading
# that as "it did not run" might well run it a second time. The wrapper warns instead, which a
# script method can only write to the host, so the preference is set rather than captured.
$WarningPreference = "SilentlyContinue"
$offlineStatement = "ALTER DATABASE [$offlineDbName] SET OFFLINE WITH ROLLBACK IMMEDIATE"

{ $offlineServer.Databases["master"].Invoke($offlineStatement) } | Should -Not -Throw

(Get-DbaDbState -SqlInstance $TestConfig.InstanceSingle -Database $offlineDbName).Status | Should -Be "OFFLINE"
}
}

Context "Databases whose names differ only in case are told apart (#10579)" -Skip:(-not $instanceIsCaseSensitive) {
BeforeAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

$caseSuffix = Get-Random
$caseDbMixed = "dbatoolsci_CaseCtx$caseSuffix"
$caseDbLower = "dbatoolsci_casectx$caseSuffix"

# The two names may differ only in case, but their files may not: NTFS is case insensitive, so
# names derived from the database name collide, first the mdf and then the ldf. So the second
# database is created under a name of its own and renamed afterwards, which keeps its files
# apart and needs no explicit file paths.
$caseDbTemporary = "dbatoolsci_casetmp$caseSuffix"
$null = New-DbaDatabase -SqlInstance $TestConfig.InstanceSingle -Name $caseDbMixed
$null = New-DbaDatabase -SqlInstance $TestConfig.InstanceSingle -Name $caseDbTemporary
$null = Invoke-DbaQuery -SqlInstance $TestConfig.InstanceSingle -Query "ALTER DATABASE [$caseDbTemporary] MODIFY NAME = [$caseDbLower]"

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

AfterAll {
$PSDefaultParameterValues["*-Dba*:EnableException"] = $true

foreach ($caseDbName in $caseDbMixed, $caseDbLower, $caseDbTemporary) {
$null = Remove-DbaDatabase -SqlInstance $TestConfig.InstanceSingle -Database $caseDbName -ErrorAction SilentlyContinue
}

$PSDefaultParameterValues.Remove("*-Dba*:EnableException")
}

It "restores the database of the caller when the two names differ only in case" {
# A case insensitive comparison reports the two as equal and skips the restore, so the caller is
# left in the wrong database - the very leak this change is about, on a valid configuration.
$caseCaller = Connect-DbaInstance -SqlInstance $TestConfig.InstanceSingle -Database $caseDbMixed -NonPooledConnection
try {
$null = $caseCaller.Databases[$caseDbLower].Query("SELECT 1")

$caseCaller.ConnectionContext.ExecuteScalar("SELECT DB_NAME()") | Should -BeExactly $caseDbMixed
} finally {
$null = $caseCaller | Disconnect-DbaInstance
}
}
}
}
87 changes: 82 additions & 5 deletions xml/dbatools.Types.ps1xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,34 @@ param (
$AllTables = $false
)

if ($AllTables) { ($this.ExecuteWithResults($Query)).Tables }
else { ($this.ExecuteWithResults($Query)).Tables[0] }
# ExecuteWithResults does not run on a private connection: the execution manager of a database is the
# connection context of the parent server, which belongs to the caller. It issues a USE and never switches
# back, so the previous database is put back here. See #10555.
$connectionContext = $this.Parent.ConnectionContext
$previousDatabase = $connectionContext.CurrentDatabase

try {
if ($AllTables) { ($this.ExecuteWithResults($Query)).Tables }
else { ($this.ExecuteWithResults($Query)).Tables[0] }
} finally {
# The comparison is case sensitive on purpose. Database names use the collation of the instance, so
# on a case sensitive instance AppDb and appdb are two different databases, and -ne would report
# them as equal and skip the restore. It cannot restore needlessly: both sides are read from the
# same property, so the same database always spells itself the same way.
if ($previousDatabase -and $connectionContext.CurrentDatabase -cne $previousDatabase) {
$escapedDatabase = $previousDatabase.Replace("]", "]]")
try {
$null = $connectionContext.ExecuteNonQuery("USE [$escapedDatabase]")
} catch {
# Putting the database back is housekeeping and must never become the outcome of the call.
# A query that makes the previous database unreachable - taking it offline, dropping it,
# renaming it, revoking access - would otherwise throw here although it succeeded, and the
# caller would read that as "it did not run" and might do it a second time. A query that
# failed would have its own error replaced by this one, which is just as misleading.
Write-Warning "The database context could not be restored to [$previousDatabase]: $($_.Exception.Message)"
}
}
}
</Script>
</ScriptMethod>
<ScriptMethod>
Expand All @@ -22,7 +48,25 @@ else { ($this.ExecuteWithResults($Query)).Tables[0] }
param (
$Command
)
$this.ExecuteNonQuery($Command)

# See the comment on Query above: this leaves the connection of the caller in this database, so the
# previous database is put back afterwards. See #10555.
$connectionContext = $this.Parent.ConnectionContext
$previousDatabase = $connectionContext.CurrentDatabase

try {
$this.ExecuteNonQuery($Command)
} finally {
# See the comment on Query above for why this is case sensitive and why a failing restore only warns.
if ($previousDatabase -and $connectionContext.CurrentDatabase -cne $previousDatabase) {
$escapedDatabase = $previousDatabase.Replace("]", "]]")
try {
$null = $connectionContext.ExecuteNonQuery("USE [$escapedDatabase]")
} catch {
Write-Warning "The database context could not be restored to [$previousDatabase]: $($_.Exception.Message)"
}
}
}
</Script>
</ScriptMethod>
</Members>
Expand All @@ -43,7 +87,24 @@ param (

try {
if ($Database) {
$dataSet = $this.Databases[$Database].ExecuteWithResults($Query)
# Going through the database object leaves the connection of the caller in that database, so the
# previous database is put back afterwards. This does not go through Database.Query, so it needs the
# same treatment of its own. See #10555.
$connectionContext = $this.ConnectionContext
$previousDatabase = $connectionContext.CurrentDatabase
try {
$dataSet = $this.Databases[$Database].ExecuteWithResults($Query)
} finally {
# See the comment on Database.Query for why this is case sensitive and why a failing restore only warns.
if ($previousDatabase -and $connectionContext.CurrentDatabase -cne $previousDatabase) {
$escapedDatabase = $previousDatabase.Replace("]", "]]")
try {
$null = $connectionContext.ExecuteNonQuery("USE [$escapedDatabase]")
} catch {
Write-Warning "The database context could not be restored to [$previousDatabase]: $($_.Exception.Message)"
}
}
}
} else {
$dataSet = $this.ConnectionContext.ExecuteWithResults($Query)
}
Expand Down Expand Up @@ -74,7 +135,23 @@ param (

try {
if ($Database) {
$this.Databases[$Database].ExecuteNonQuery($Command)
# See the comment on Query above: this does not go through Database.Invoke either, so the previous
# database is put back here as well. See #10555.
$connectionContext = $this.ConnectionContext
$previousDatabase = $connectionContext.CurrentDatabase
try {
$this.Databases[$Database].ExecuteNonQuery($Command)
} finally {
# See the comment on Database.Query for why this is case sensitive and why a failing restore only warns.
if ($previousDatabase -and $connectionContext.CurrentDatabase -cne $previousDatabase) {
$escapedDatabase = $previousDatabase.Replace("]", "]]")
try {
$null = $connectionContext.ExecuteNonQuery("USE [$escapedDatabase]")
} catch {
Write-Warning "The database context could not be restored to [$previousDatabase]: $($_.Exception.Message)"
}
}
}
} else {
$this.ConnectionContext.ExecuteNonQuery($Command)
}
Expand Down