Restore script update

RestoreScript.ps1

@@ -19,7 +19,16 @@ param(
     [string]$LogPath,
 
     [Parameter(Mandatory=$false)]
-    [string]$NewName
+    [string]$NewName,
+
+    [Parameter(Mandatory=$false)]
+    [string]$StopAtTime,
+
+    [Parameter(Mandatory=$false)]
+    [string]$StopAtLSN,
+
+    [Parameter(Mandatory=$false)]
+    [switch]$PreviewRestorePlan
 )
 
 # Function to catalog backups
@@ -98,9 +107,23 @@ function Report-Catalog {
 
 # Function to restore database
 function Restore-Database {
-    param([string]$DbName, [hashtable]$Catalog, [string]$Instance, [string]$DataPath, [string]$LogPath, [string]$TargetDbName)
+    param(
+        [string]$DbName,
+        [hashtable]$Catalog,
+        [string]$Instance,
+        [string]$DataPath,
+        [string]$LogPath,
+        [string]$TargetDbName,
+        [string]$StopAtTime,
+        [string]$StopAtLSN,
+        [bool]$PreviewRestorePlan
+    )
 
     $restoreDbName = if ($TargetDbName) { $TargetDbName } else { $DbName }
+    $shouldRenamePhysicalFiles = -not [string]::IsNullOrWhiteSpace($TargetDbName) -and ($TargetDbName -ne $DbName)
+    $isPointInTimeRestore = ($StopAtTime -or $StopAtLSN)
+    $targetStopAtTime = $null
+    $targetStopAtLsn = $null
 
     if (-not $catalog.ContainsKey($DbName)) {
         Write-Error "Database $DbName not found in catalog"
@@ -115,14 +138,75 @@ function Restore-Database {
         return
     }
 
-    $latestFullKey = $dbCatalog['FULL'].Keys | Sort-Object -Descending | Select-Object -First 1
+    if ($StopAtTime) {
-    $fullFiles = $dbCatalog['FULL'][$latestFullKey] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
+        try {
-    $fullHeader = Get-BackupInfo -Files $fullFiles -Instance $Instance
+            $targetStopAtTime = [datetime]::Parse($StopAtTime)
-    if (-not $fullHeader) {
+        } catch {
-        Write-Error "Failed to read FULL backup header for database $DbName"
+            Write-Error "StopAtTime '$StopAtTime' is not a valid datetime value"
+            return
+        }
+    }
+
+    if ($StopAtLSN) {
+        $lsnForCompare = if ($StopAtLSN -like 'lsn:*') { $StopAtLSN.Substring(4) } else { $StopAtLSN }
+        try {
+            $targetStopAtLsn = [decimal]$lsnForCompare
+        } catch {
+            Write-Error "StopAtLSN '$StopAtLSN' must be a numeric LSN value"
+            return
+        }
+    }
+
+    # Select FULL backup. For StopAtTime, choose latest FULL that finished at/before target time.
+    $fullCandidates = @()
+    $fullKeys = $dbCatalog['FULL'].Keys | Sort-Object
+    foreach ($key in $fullKeys) {
+        $files = $dbCatalog['FULL'][$key] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
+        $header = Get-BackupInfo -Files $files -Instance $Instance
+        if (-not $header) {
+            Write-Warning "Skipping FULL backup $key for $DbName because its header could not be read"
+            continue
+        }
+
+        $finish = $null
+        try {
+            $finish = [datetime]$header.BackupFinishDate
+        } catch {
+            Write-Warning "Skipping FULL backup $key for $DbName because BackupFinishDate could not be parsed"
+            continue
+        }
+
+        if ($targetStopAtTime -and $finish -gt $targetStopAtTime) {
+            continue
+        }
+
+        $fullCandidates += @{
+            Key = $key
+            Files = $files
+            Header = $header
+            Finish = $finish
+        }
+    }
+
+    if ($fullCandidates.Count -eq 0) {
+        if ($targetStopAtTime) {
+            Write-Error "No FULL backup found that is at or before StopAtTime '$StopAtTime' for database $DbName"
+        } else {
+            Write-Error "No usable FULL backup found for database $DbName"
+        }
         return
     }
 
+    $selectedFull = if ($targetStopAtTime) {
+        $fullCandidates | Sort-Object { $_.Finish } -Descending | Select-Object -First 1
+    } else {
+        $fullCandidates | Sort-Object { $_.Key } -Descending | Select-Object -First 1
+    }
+
+    $latestFullKey = $selectedFull.Key
+    $fullFiles = $selectedFull.Files
+    $fullHeader = $selectedFull.Header
+
     $fullBaseLsn = $fullHeader.CheckpointLSN
     if (-not $fullBaseLsn) {
         $fullBaseLsn = $fullHeader.FirstLSN
@@ -153,7 +237,7 @@ function Restore-Database {
 
         # When restoring as a different database name, rewrite physical file names
         # so they do not clash with the source database files on disk.
-        if ($TargetDbName) {
+        if ($shouldRenamePhysicalFiles) {
             if ($type -eq 'D') {
                 $dataFileIndex++
                 $newBaseName = if ($dataFileIndex -eq 1) { "${restoreDbName}_Data" } else { "${restoreDbName}_Data$dataFileIndex" }
@@ -174,17 +258,13 @@ function Restore-Database {
         }
     }
 
-    # Restore FULL with NORECOVERY
+    # Build FULL restore query
     $withClause = if ($moveClauses) { "WITH $($moveClauses -join ', '), NORECOVERY" } else { "WITH NORECOVERY" }
     $restoreQuery = "RESTORE DATABASE [$restoreDbName] FROM $($fileListStr -join ', ') $withClause"
-    if ($TargetDbName) {
+    $fullRestoreQuery = $restoreQuery
-        Write-Host "Restoring FULL backup for $DbName as $restoreDbName..."
-    } else {
-        Write-Host "Restoring FULL backup for $DbName..."
-    }
-    Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop
 
     # Choose the latest DIFF that is compatible with the selected FULL.
+    $selectedDiff = $null
     $selectedDiffKey = $null
     if ($dbCatalog.ContainsKey('DIFF')) {
         $compatibleDiffs = @()
@@ -197,30 +277,212 @@ function Restore-Database {
                 continue
             }
 
-            if ($diffHeader.DifferentialBaseLSN -eq $fullBaseLsn) {
+            if ($diffHeader.DifferentialBaseLSN -ne $fullBaseLsn) {
-                $compatibleDiffs += @{ Key = $key; Files = $diffFiles }
+                continue
             }
+
+            $diffFinish = $null
+            try {
+                $diffFinish = [datetime]$diffHeader.BackupFinishDate
+            } catch {
+                Write-Warning "Skipping DIFF backup $key for $DbName because BackupFinishDate could not be parsed"
+                continue
+            }
+
+            if ($targetStopAtTime -and $diffFinish -gt $targetStopAtTime) {
+                continue
+            }
+
+            $compatibleDiffs += @{ Key = $key; Files = $diffFiles; Finish = $diffFinish }
         }
 
         if ($compatibleDiffs.Count -gt 0) {
-            $selectedDiff = $compatibleDiffs | Sort-Object { $_.Key } -Descending | Select-Object -First 1
+            $selectedDiff = if ($targetStopAtTime) {
+                $compatibleDiffs | Sort-Object { $_.Finish } -Descending | Select-Object -First 1
+            } else {
+                $compatibleDiffs | Sort-Object { $_.Key } -Descending | Select-Object -First 1
+            }
             $selectedDiffKey = $selectedDiff.Key
-            $fileList = $selectedDiff.Files | ForEach-Object { "DISK = '$($_.FullName)'" }
-            $restoreQuery = "RESTORE DATABASE [$restoreDbName] FROM $($fileList -join ', ') WITH NORECOVERY"
-            Write-Host "Applying DIFF backup $selectedDiffKey for $restoreDbName..."
-            Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop
         }
     }
 
     $logStartKey = if ($selectedDiffKey) { $selectedDiffKey } else { $latestFullKey }
+    $logKeys = @()
 
-    # Apply LOG backups after the latest restored FULL/DIFF point
+    # Select LOG backups after the latest restored FULL/DIFF point.
     if ($dbCatalog.ContainsKey('LOG')) {
-        $logKeys = $dbCatalog['LOG'].Keys | Where-Object { $_ -gt $logStartKey } | Sort-Object
+        $candidateLogKeys = @($dbCatalog['LOG'].Keys | Where-Object { $_ -gt $logStartKey } | Sort-Object)
-        foreach ($key in $logKeys) {
+        if ($isPointInTimeRestore -and $candidateLogKeys.Count -eq 0) {
+            Write-Error "Point-in-time restore requested but no LOG backups are available after the selected FULL/DIFF restore point"
+            return
+        }
+
+        if ($StopAtTime) {
+            $cutoffFound = $false
+            foreach ($key in $candidateLogKeys) {
+                $logFiles = $dbCatalog['LOG'][$key] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
+                $logHeader = Get-BackupInfo -Files $logFiles -Instance $Instance
+                if (-not $logHeader) {
+                    continue
+                }
+
+                $logStartTime = $null
+                $logFinishTime = $null
+                try {
+                    $logStartTime = [datetime]$logHeader.BackupStartDate
+                    $logFinishTime = [datetime]$logHeader.BackupFinishDate
+                } catch {
+                    continue
+                }
+
+                if ($logFinishTime -lt $targetStopAtTime) {
+                    $logKeys += $key
+                    continue
+                }
+
+                # A log backup contains all records up to BackupFinishDate. The first log
+                # whose finish time meets/exceeds StopAtTime is the PITR cutoff log.
+                if ($logFinishTime -ge $targetStopAtTime) {
+                    $logKeys += $key
+                    $cutoffFound = $true
+                    break
+                }
+            }
+
+            if (-not $cutoffFound) {
+                Write-Error "StopAtTime '$StopAtTime' is not covered by available LOG backups after the selected FULL/DIFF restore point"
+                return
+            }
+        } elseif ($StopAtLSN) {
+            $cutoffFound = $false
+            foreach ($key in $candidateLogKeys) {
+                $logFiles = $dbCatalog['LOG'][$key] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
+                $logHeader = Get-BackupInfo -Files $logFiles -Instance $Instance
+                if (-not $logHeader) {
+                    continue
+                }
+
+                $logFirstLsn = $null
+                $logLastLsn = $null
+                try {
+                    $logFirstLsn = [decimal]$logHeader.FirstLSN
+                    $logLastLsn = [decimal]$logHeader.LastLSN
+                } catch {
+                    continue
+                }
+
+                if ($logLastLsn -lt $targetStopAtLsn) {
+                    $logKeys += $key
+                    continue
+                }
+
+                if ($logFirstLsn -le $targetStopAtLsn -and $targetStopAtLsn -le $logLastLsn) {
+                    $logKeys += $key
+                    $cutoffFound = $true
+                    break
+                }
+
+                if ($logFirstLsn -gt $targetStopAtLsn) {
+                    break
+                }
+            }
+
+            if (-not $cutoffFound) {
+                Write-Error "StopAtLSN '$StopAtLSN' is not covered by available LOG backups after the selected FULL/DIFF restore point"
+                return
+            }
+        } else {
+            $logKeys = $candidateLogKeys
+        }
+
+        if ($isPointInTimeRestore -and $logKeys.Count -eq 0) {
+            Write-Error "Point-in-time restore requested but no usable LOG backups were selected"
+            return
+        }
+    } elseif ($isPointInTimeRestore) {
+        Write-Error "Point-in-time restore requested but this database has no LOG backups in the catalog"
+        return
+    }
+
+    if ($PreviewRestorePlan) {
+        Write-Host ""
+        Write-Host "Restore Plan Preview"
+        Write-Host "===================="
+        Write-Host "Source Database: $DbName"
+        Write-Host "Target Database: $restoreDbName"
+        Write-Host "FULL Backup Key: $latestFullKey"
+        Write-Host "FULL Backup Files:"
+        foreach ($f in $fullFiles) {
+            Write-Host "  - $($f.FullName)"
+        }
+        if ($selectedDiffKey) {
+            Write-Host "Selected DIFF Key: $selectedDiffKey"
+            Write-Host "Selected DIFF Files:"
+            foreach ($f in $selectedDiff.Files) {
+                Write-Host "  - $($f.FullName)"
+            }
+        } else {
+            Write-Host "Selected DIFF Key: none"
+        }
+        if ($logKeys.Count -gt 0) {
+            Write-Host "LOG Range: $($logKeys[0]) -> $($logKeys[-1]) ($($logKeys.Count) backup sets)"
+        } else {
+            Write-Host "LOG Range: none"
+        }
+        if ($isPointInTimeRestore) {
+            if ($StopAtTime) {
+                Write-Host "PITR Target: STOPAT '$StopAtTime'"
+            } else {
+                $previewLsn = if ($StopAtLSN -like 'lsn:*') { $StopAtLSN } else { "lsn:$StopAtLSN" }
+                Write-Host "PITR Target: STOPATMARK '$previewLsn'"
+            }
+        } else {
+            Write-Host "PITR Target: none (restore to latest available point)"
+        }
+        Write-Host "MOVE Clauses:"
+        foreach ($m in $moveClauses) {
+            Write-Host "  - $m"
+        }
+        Write-Host ""
+
+        Write-Host "PreviewRestorePlan enabled. Exiting without executing restore operations."
+        return
+    }
+
+    if ($shouldRenamePhysicalFiles) {
+        Write-Host "Restoring FULL backup for $DbName as $restoreDbName..."
+    } else {
+        Write-Host "Restoring FULL backup for $DbName..."
+    }
+    Invoke-Sqlcmd -ServerInstance $Instance -Query $fullRestoreQuery -QueryTimeout 0 -ErrorAction Stop
+
+    if ($selectedDiffKey) {
+        $fileList = $selectedDiff.Files | ForEach-Object { "DISK = '$($_.FullName)'" }
+        $restoreQuery = "RESTORE DATABASE [$restoreDbName] FROM $($fileList -join ', ') WITH NORECOVERY"
+        Write-Host "Applying DIFF backup $selectedDiffKey for $restoreDbName..."
+        Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop
+    }
+
+    if ($logKeys.Count -gt 0) {
+        for ($i = 0; $i -lt $logKeys.Count; $i++) {
+            $key = $logKeys[$i]
             $logFiles = $dbCatalog['LOG'][$key] | Sort-Object { $_.Stripe } | ForEach-Object { $_.File }
             $fileList = $logFiles | ForEach-Object { "DISK = '$($_.FullName)'" }
-            $restoreQuery = "RESTORE LOG [$restoreDbName] FROM $($fileList -join ', ') WITH NORECOVERY"
+            $logWithClause = "WITH NORECOVERY"
+
+            # Apply STOPAT/STOPATMARK only on the final LOG restore statement.
+            if ($isPointInTimeRestore -and $i -eq ($logKeys.Count - 1)) {
+                if ($StopAtTime) {
+                    $escapedStopAtTime = $StopAtTime.Replace("'", "''")
+                    $logWithClause = "$logWithClause, STOPAT = '$escapedStopAtTime'"
+                } elseif ($StopAtLSN) {
+                    $normalizedLsn = if ($StopAtLSN -like "lsn:*") { $StopAtLSN } else { "lsn:$StopAtLSN" }
+                    $escapedLsn = $normalizedLsn.Replace("'", "''")
+                    $logWithClause = "$logWithClause, STOPATMARK = '$escapedLsn'"
+                }
+            }
+
+            $restoreQuery = "RESTORE LOG [$restoreDbName] FROM $($fileList -join ', ') $logWithClause"
             Write-Host "Applying LOG backup $key for $restoreDbName..."
             Invoke-Sqlcmd -ServerInstance $Instance -Query $restoreQuery -QueryTimeout 0 -ErrorAction Stop
         }
@@ -244,8 +506,10 @@ function Print-BackupSummary {
     # FULL backups
     if ($Headers.ContainsKey('FULL')) {
         Write-Host "FULL Backups:"
-        foreach ($item in $Headers['FULL'] | Sort-Object { $_.Key }) {
+        $fullItems = @($Headers['FULL'] | Sort-Object { $_.Key })
-            $header = $item.Header
+        foreach ($item in $fullItems) {
+            $header = @($item.Header)[0]
+            if (-not $header) { continue }
             Write-Host "  Date: $($header.BackupFinishDate) | LSN Range: $($header.FirstLSN) - $($header.LastLSN)"
         }
     }
@@ -253,18 +517,20 @@ function Print-BackupSummary {
     # DIFF backups
     if ($Headers.ContainsKey('DIFF')) {
         Write-Host "DIFFERENTIAL Backups:"
-        foreach ($item in $Headers['DIFF'] | Sort-Object { $_.Key }) {
+        $diffItems = @($Headers['DIFF'] | Sort-Object { $_.Key })
-            $header = $item.Header
+        foreach ($item in $diffItems) {
+            $header = @($item.Header)[0]
+            if (-not $header) { continue }
             Write-Host "  Date: $($header.BackupFinishDate) | Base LSN: $($header.DifferentialBaseLSN) | LSN Range: $($header.FirstLSN) - $($header.LastLSN)"
         }
     }
 
     # LOG backups
     if ($Headers.ContainsKey('LOG')) {
-        $logItems = $Headers['LOG'] | Sort-Object { $_.Key }
+        $logItems = @($Headers['LOG'] | Sort-Object { $_.Key })
         if ($logItems.Count -gt 0) {
-            $firstLog = $logItems[0].Header
+            $firstLog = @($logItems[0].Header)[0]
-            $lastLog = $logItems[-1].Header
+            $lastLog = @($logItems[-1].Header)[0]
             Write-Host "LOG Backups:"
             Write-Host "  Point-in-Time Range: $($firstLog.BackupStartDate) to $($lastLog.BackupFinishDate)"
             Write-Host "  LSN Range: $($firstLog.FirstLSN) - $($lastLog.LastLSN)"
@@ -272,8 +538,11 @@ function Print-BackupSummary {
             # Check for gaps
             $gaps = @()
             for ($i = 1; $i -lt $logItems.Count; $i++) {
-                $prevLast = $logItems[$i-1].Header.LastLSN
+                $prevHeader = @($logItems[$i-1].Header)[0]
-                $currFirst = $logItems[$i].Header.FirstLSN
+                $currHeader = @($logItems[$i].Header)[0]
+                if (-not $prevHeader -or -not $currHeader) { continue }
+                $prevLast = $prevHeader.LastLSN
+                $currFirst = $currHeader.FirstLSN
                 if ($prevLast -ne $currFirst) {
                     $gaps += "Gap between $($logItems[$i-1].Key) (LSN $($prevLast)) and $($logItems[$i].Key) (LSN $($currFirst))"
                 }
@@ -368,6 +637,21 @@ if ($Action -ne "restore" -and $NewName) {
     exit 1
 }
 
+if ($StopAtTime -and $StopAtLSN) {
+    Write-Error "Specify only one of StopAtTime or StopAtLSN"
+    exit 1
+}
+
+if ($Action -ne "restore" -and ($StopAtTime -or $StopAtLSN)) {
+    Write-Error "StopAtTime and StopAtLSN are only valid with restore action"
+    exit 1
+}
+
+if ($Action -ne "restore" -and $PreviewRestorePlan) {
+    Write-Error "PreviewRestorePlan is only valid with restore action"
+    exit 1
+}
+
 if ($Action -eq "catalog") {
     $catalog = Catalog-Backups -Path $LiveMountRoot
     if ($DatabaseName) {
@@ -390,7 +674,7 @@ if ($Action -eq "catalog") {
     }
 
     $catalog = Catalog-Backups -Path $LiveMountRoot
-    Restore-Database -DbName $DatabaseName -Catalog $catalog -Instance $SqlInstance -DataPath $DataPath -LogPath $LogPath -TargetDbName $NewName
+    Restore-Database -DbName $DatabaseName -Catalog $catalog -Instance $SqlInstance -DataPath $DataPath -LogPath $LogPath -TargetDbName $NewName -StopAtTime $StopAtTime -StopAtLSN $StopAtLSN -PreviewRestorePlan $PreviewRestorePlan
 } elseif ($Action -eq "verify") {
     $catalog = Catalog-Backups -Path $LiveMountRoot
     if ($DatabaseName) {