Invoke-DbaDbShrink - Add WAIT_AT_LOW_PRIORITY support

public/Invoke-DbaDbShrink.ps1

@@ -58,6 +58,16 @@ function Invoke-DbaDbShrink {
         Sets the command timeout in minutes for the shrink operation. Defaults to 0 (infinite timeout).
         Large database shrinks can take hours to complete, so the default allows operations to run without timing out.
 
+    .PARAMETER WaitAtLowPriority
+        Instructs DBCC SHRINKFILE to wait at low priority for schema modification locks, minimizing blocking of other sessions.
+        Requires SQL Server 2022 (version 16) or later. When a lock cannot be obtained, SQL Server will abort the operation
+        after approximately one minute and log error 49516. The AbortAfterWait parameter controls what is aborted.
+
+    .PARAMETER AbortAfterWait
+        Specifies what to abort when the WAIT_AT_LOW_PRIORITY wait period expires. Valid values are Self and Blockers.
+        Self aborts the shrink operation itself. Blockers kills the user sessions that block the lock acquisition.
+        Defaults to Self. Only applies when WaitAtLowPriority is specified.
+
     .PARAMETER LogsOnly
         Deprecated. Use FileType instead.
 
@@ -165,6 +175,11 @@ function Invoke-DbaDbShrink {
 
         Shrinks all databases coming from a pre-filtered list via Get-DbaDatabase
 
+    .EXAMPLE
+        PS C:\> Invoke-DbaDbShrink -SqlInstance sql2022 -Database AdventureWorks2014 -WaitAtLowPriority -AbortAfterWait Blockers
+
+        Shrinks AdventureWorks2014 using WAIT_AT_LOW_PRIORITY, killing blocking sessions if the wait expires. Requires SQL Server 2022+.
+
     #>
     [CmdletBinding(SupportsShouldProcess, ConfirmImpact = 'Low')]
     param (
@@ -183,6 +198,9 @@ function Invoke-DbaDbShrink {
         [string]$FileType = 'All',
         [int64]$StepSize,
         [int]$StatementTimeout = 0,
+        [switch]$WaitAtLowPriority,
+        [ValidateSet('Self', 'Blockers')]
+        [string]$AbortAfterWait = 'Self',
         [switch]$ExcludeIndexStats,
         [switch]$ExcludeUpdateUsage,
         [switch]$EnableException,
@@ -202,6 +220,11 @@ function Invoke-DbaDbShrink {
         }
         $StatementTimeoutSeconds = $StatementTimeout * 60
 
+        $walp = ""
+        if ($WaitAtLowPriority) {
+            $walp = " WITH WAIT_AT_LOW_PRIORITY (ABORT_AFTER_WAIT = $($AbortAfterWait.ToUpper()))"
+        }
+
         $sql = 'SELECT
                   AVG(avg_fragmentation_in_percent) AS [avg_fragmentation_in_percent]
                 , MAX(avg_fragmentation_in_percent) AS [max_fragmentation_in_percent]
@@ -255,6 +278,10 @@ function Invoke-DbaDbShrink {
                 continue
             }
 
+            if ($WaitAtLowPriority -and $instance.VersionMajor -lt 16) {
+                Stop-Function -Message "WAIT_AT_LOW_PRIORITY for DBCC SHRINKFILE requires SQL Server 2022 (version 16) or later. $instance is running version $($instance.VersionMajor)." -Target $instance -Continue
+            }
+
             $files = @()
             if ($FileType -in ('Log', 'All')) {
                 $files += $db.LogFiles
@@ -279,6 +306,8 @@ function Invoke-DbaDbShrink {
                 Write-Message -Level Verbose -Message "Target Freespace: $($desiredSpaceAvailableKB)"
                 Write-Message -Level Verbose -Message "Target FileSize: $($desiredFileSizeKB)"
 
+                $escapedFileName = $file.Name.Replace("'", "''")
+
                 if ($spaceAvailableKB -le $desiredSpaceAvailableKB) {
                     Write-Message -Level Warning -Message "File size of ($startingSizeKB) is less than or equal to the desired outcome ($desiredFileSizeKB) for $($file.Name)"
                 } else {
@@ -313,7 +342,14 @@ function Invoke-DbaDbShrink {
                                         $shrinkSizeKB = $desiredFileSizeKB
                                     }
                                     Write-Message -Level Verbose -Message ('Shrinking {0} to {1}' -f $file.Name, $shrinkSizeKB)
-                                    $file.Shrink($shrinkSizeKB.Megabyte, $ShrinkMethod)
+                                    $targetMB = [int]$shrinkSizeKB.Megabyte
+                                    $shrinkSqlArgs = switch ($ShrinkMethod) {
+                                        'EmptyFile'    { "N'$escapedFileName', EMPTYFILE" }
+                                        'NoTruncate'   { "N'$escapedFileName', $targetMB, NOTRUNCATE" }
+                                        'TruncateOnly' { "N'$escapedFileName', $targetMB, TRUNCATEONLY" }
+                                        default        { "N'$escapedFileName', $targetMB" }
+                                    }
+                                    $null = $instance.Query("DBCC SHRINKFILE ($shrinkSqlArgs)$walp", $db.name)
                                     $file.Refresh()
 
                                     if ($startingSizeKB -eq ($file.Size * 1024)) {
@@ -322,7 +358,14 @@ function Invoke-DbaDbShrink {
                                     }
                                 }
                             } else {
-                                $file.Shrink(($desiredFileSizeKB.Megabyte), $ShrinkMethod)
+                                $targetMB = [int]$desiredFileSizeKB.Megabyte
+                                $shrinkSqlArgs = switch ($ShrinkMethod) {
+                                    'EmptyFile'    { "N'$escapedFileName', EMPTYFILE" }
+                                    'NoTruncate'   { "N'$escapedFileName', $targetMB, NOTRUNCATE" }
+                                    'TruncateOnly' { "N'$escapedFileName', $targetMB, TRUNCATEONLY" }
+                                    default        { "N'$escapedFileName', $targetMB" }
+                                }
+                                $null = $instance.Query("DBCC SHRINKFILE ($shrinkSqlArgs)$walp", $db.name)
                                 $file.Refresh()
                             }
                             $success = $true

tests/Invoke-DbaDbShrink.Tests.ps1

@@ -21,6 +21,8 @@ Describe $CommandName -Tag UnitTests {
                 "FileType",
                 "StepSize",
                 "StatementTimeout",
+                "WaitAtLowPriority",
+                "AbortAfterWait",
                 "ExcludeIndexStats",
                 "ExcludeUpdateUsage",
                 "EnableException",
@@ -155,5 +157,78 @@ Describe $CommandName -Tag IntegrationTests {
             $db.FileGroups[0].Files[0].Size | Should -BeLessThan $oldDataSize
             $db.LogFiles[0].Size | Should -Be $oldLogSize
         }
+
+        It "Shrinks with WaitAtLowPriority on SQL Server 2022+" {
+            if ($server.VersionMajor -lt 16) {
+                Set-ItResult -Skipped -Because "Test is only for SQL Server 2022 and later"
+                return
+            }
+
+            $result = Invoke-DbaDbShrink $server -Database $db.Name -FileType Data -WaitAtLowPriority -AbortAfterWait Self
+            $result.Database | Should -Be $db.Name
+            $result.File | Should -Be $db.Name
+            $result.Success | Should -Be $true
+            $db.Refresh()
+            $db.RecalculateSpaceUsage()
+            $db.FileGroups[0].Files[0].Refresh()
+            $db.FileGroups[0].Files[0].Size | Should -BeLessThan $oldDataSize
+        }
+
+        It "Verifies WAIT_AT_LOW_PRIORITY SQL syntax appears in running requests when blocked" {
+            if ($server.VersionMajor -lt 16) {
+                Set-ItResult -Skipped -Because "Test is only for SQL Server 2022 and later"
+                return
+            }
+
+            # Create a table with data so DBCC SHRINKFILE will need to acquire locks when moving pages
+            $null = $server.Query("USE [$($db.Name)]; CREATE TABLE dbo.dbatoolsci_shrink_blocker (c1 INT); INSERT INTO dbo.dbatoolsci_shrink_blocker SELECT TOP 1000 ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) FROM sys.all_columns a1 CROSS JOIN sys.all_columns a2")
+
+            # Start a job that holds an exclusive table lock to force the shrink to wait at low priority
+            $blockConnStr = $server.ConnectionContext.ConnectionString
+            $blockDbName = $db.Name
+            $blockJob = Start-Job -ScriptBlock {
+                param($connStr, $dbName)
+                $conn = New-Object System.Data.SqlClient.SqlConnection($connStr)
+                $conn.Open()
+                $cmd = $conn.CreateCommand()
+                $cmd.CommandTimeout = 60
+                $cmd.CommandText = "USE [$dbName]; BEGIN TRAN; SELECT TOP 1 c1 FROM dbo.dbatoolsci_shrink_blocker WITH (TABLOCKX, HOLDLOCK); WAITFOR DELAY '00:00:20'; IF @@TRANCOUNT > 0 ROLLBACK TRAN"
+                try { $cmd.ExecuteNonQuery() } catch { }
+                $conn.Close()
+            } -ArgumentList $blockConnStr, $blockDbName
+
+            Start-Sleep -Seconds 2
+
+            # Start the shrink with WAIT_AT_LOW_PRIORITY in a background job
+            $shrinkServerName = $server.DomainInstanceName
+            $shrinkDbName = $db.Name
+            $shrinkModulePath = (Get-Module dbatools | Select-Object -First 1).Path
+            $shrinkJob = Start-Job -ScriptBlock {
+                param($modulePath, $serverName, $dbName)
+                Import-Module $modulePath
+                Invoke-DbaDbShrink -SqlInstance $serverName -Database $dbName -FileType Data -WaitAtLowPriority
+            } -ArgumentList $shrinkModulePath, $shrinkServerName, $shrinkDbName
+
+            Start-Sleep -Seconds 3
+
+            # Verify the SQL text of the running shrink request contains WAIT_AT_LOW_PRIORITY
+            $sqlTextCount = ($server.Query("SELECT COUNT(*) AS C FROM sys.dm_exec_requests r CROSS APPLY sys.dm_exec_sql_text(r.sql_handle) t WHERE t.text LIKE '%WAIT_AT_LOW_PRIORITY%'")).C
+
+            $null = $blockJob | Wait-Job -Timeout 30
+            $blockJob | Remove-Job -Force
+            $null = $shrinkJob | Wait-Job -Timeout 60
+            $shrinkJob | Remove-Job -Force
+
+            $sqlTextCount | Should -BeGreaterThan 0
+        }
+
+        It "Returns an error when WaitAtLowPriority is used on SQL Server older than 2022" {
+            if ($server.VersionMajor -ge 16) {
+                Set-ItResult -Skipped -Because "Test is only for SQL Server 2019 and older"
+                return
+            }
+            $result = Invoke-DbaDbShrink $server -Database $db.Name -FileType Data -WaitAtLowPriority -WarningAction SilentlyContinue
+            $WarnVar | Should -Match "SQL Server 2022"
+        }
     }
 }