SqlAgDatabase: Fixes Issue #1358 (#1372)

- Changes to SqlServerDsc.Common - Added StatementTimeout to function 'Connect-SQL' with default 600 seconds (10mins). - Added StatementTimeout to function 'Invoke-Query' with default 600 seconds (10mins). Fixes Issue#1358 - Changes to SqlAGDatabase - Added new parameter 'ReplaceExisting' with default false. This allows forced restores when a database already exists on secondary. - Added StatementTimeout to Invoke-Query to fix Issue#1358
dsccommunity · Jun 6, 2019 · e62c1f5 · e62c1f5
1 parent b2c5778
commit e62c1f5
Show file tree

Hide file tree

Showing 7 changed files with 117 additions and 23 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,14 @@
 
 ## Unreleased
 
+- Changes to SqlServerDsc.Common
+  - Added StatementTimeout to function 'Connect-SQL' with default 600 seconds (10mins).
+  - Added StatementTimeout to function 'Invoke-Query' with default 600 seconds (10mins).
+    Fixes Issue#1358
+- Changes to SqlAGDatabase
+  - Added new parameter 'ReplaceExisting' with default false.
+    This allows forced restores when a database already exists on secondary.
+  - Added StatementTimeout to Invoke-Query to fix Issue#1358
 - Changes to SqlServerDsc
   - Opt-in to the common test 'Common Test - Validation Localization'.
   - Opt-in to the common test 'Common Test - Flagged Script Analyzer Rules'

diff --git a/DSCResources/MSFT_SqlAGDatabase/MSFT_SqlAGDatabase.psm1 b/DSCResources/MSFT_SqlAGDatabase/MSFT_SqlAGDatabase.psm1
@@ -60,13 +60,14 @@ function Get-TargetResource
     # Create an object that reflects the current configuration
     $currentConfiguration = @{
         DatabaseName          = @()
-        ServerName             = $ServerName
-        InstanceName       = $InstanceName
+        ServerName            = $ServerName
+        InstanceName          = $InstanceName
         AvailabilityGroupName = ''
         BackupPath            = ''
         Ensure                = ''
         Force                 = $false
         MatchDatabaseOwner    = $false
+        ReplaceExisting       = $false
         IsActiveNode          = $false
     }
 
@@ -144,6 +145,11 @@ function Get-TargetResource
 
         The default is '$false'.
 
+    .PARAMETER ReplaceExisting
+        If set to $true, this adds the restore option WITH REPLACE.
+        If set to $false, Existing databases and files will block the restore and throw error.
+        The default is '$false'.
+
     .PARAMETER ProcessOnlyOnActiveNode
         Specifies that the resource will only determine if a change is needed if the target node is the active host of the SQL Server Instance.
         Not used in Set-TargetResource.
@@ -186,6 +192,10 @@ function Set-TargetResource
         [System.Boolean]
         $MatchDatabaseOwner,
 
+        [Parameter()]
+        [System.Boolean]
+        $ReplaceExisting,
+
         [Parameter()]
         [System.Boolean]
         $ProcessOnlyOnActiveNode
@@ -415,7 +425,7 @@ function Set-TargetResource
                     ErrorAction    = 'Stop'
                 }
 
-                # If no full backup was ever taken, do not take a backup with CopyOnly
+                # If database object last backup data not equal to 0 then backup with CopyOnly.
                 if ( $databaseObject.LastBackupDate -ne 0 )
                 {
                     $backupSqlDatabaseParameters.Add('CopyOnly', $true)
@@ -486,6 +496,12 @@ function Set-TargetResource
                 $restoreDatabaseQueryStringBuilder.Append($databaseFullBackupFile) | Out-Null
                 $restoreDatabaseQueryStringBuilder.AppendLine('''') | Out-Null
                 $restoreDatabaseQueryStringBuilder.Append('WITH NORECOVERY') | Out-Null
+
+                if ( $ReplaceExisting )
+                {
+                    $restoreDatabaseQueryStringBuilder.Append(',REPLACE') | Out-Null
+                }
+
                 if ( $MatchDatabaseOwner )
                 {
                     $restoreDatabaseQueryStringBuilder.AppendLine() | Out-Null
@@ -510,11 +526,13 @@ function Set-TargetResource
                 $restoreLogQueryStringBuilder.Append($databaseLogBackupFile) | Out-Null
                 $restoreLogQueryStringBuilder.AppendLine('''') | Out-Null
                 $restoreLogQueryStringBuilder.Append('WITH NORECOVERY') | Out-Null
+
                 if ( $MatchDatabaseOwner )
                 {
                     $restoreLogQueryStringBuilder.AppendLine() | Out-Null
                     $restoreLogQueryStringBuilder.Append('REVERT') | Out-Null
                 }
+
                 $restoreLogQueryString = $restoreLogQueryStringBuilder.ToString()
 
                 try
@@ -527,8 +545,8 @@ function Set-TargetResource
                         $currentReplicaAvailabilityGroupObject = $currentAvailabilityGroupReplicaServerObject.AvailabilityGroups[$AvailabilityGroupName]
 
                         # Restore the database
-                        Invoke-Query -SQLServer $currentAvailabilityGroupReplicaServerObject.NetName -SQLInstanceName $currentAvailabilityGroupReplicaServerObject.ServiceName -Database master -Query $restoreDatabaseQueryString
-                        Invoke-Query -SQLServer $currentAvailabilityGroupReplicaServerObject.NetName -SQLInstanceName $currentAvailabilityGroupReplicaServerObject.ServiceName -Database master -Query $restoreLogQueryString
+                        Invoke-Query -SQLServer $currentAvailabilityGroupReplicaServerObject.NetName -SQLInstanceName $currentAvailabilityGroupReplicaServerObject.ServiceName -Database master -Query $restoreDatabaseQueryString -StatementTimeout 0
+                        Invoke-Query -SQLServer $currentAvailabilityGroupReplicaServerObject.NetName -SQLInstanceName $currentAvailabilityGroupReplicaServerObject.ServiceName -Database master -Query $restoreLogQueryString -StatementTimeout 0
 
                         # Add the database to the Availability Group
                         Add-SqlAvailabilityDatabase -InputObject $currentReplicaAvailabilityGroupObject -Database $databaseToAddToAvailabilityGroup
@@ -637,6 +655,11 @@ function Set-TargetResource
 
         The default is '$false'.
 
+    .PARAMETER ReplaceExisting
+        If set to $true, this adds the restore option WITH REPLACE.
+        If set to $false, Existing databases and files will block the restore and throw error.
+        The default is '$false'.
+
     .PARAMETER ProcessOnlyOnActiveNode
         Specifies that the resource will only determine if a change is needed if the target node is the active host of the SQL Server Instance.
 #>
@@ -679,6 +702,10 @@ function Test-TargetResource
         [System.Boolean]
         $MatchDatabaseOwner,
 
+        [Parameter()]
+        [System.Boolean]
+        $ReplaceExisting,
+
         [Parameter()]
         [System.Boolean]
         $ProcessOnlyOnActiveNode
@@ -688,8 +715,8 @@ function Test-TargetResource
 
     $getTargetResourceParameters = @{
         DatabaseName          = $DatabaseName
-        ServerName             = $ServerName
-        InstanceName       = $InstanceName
+        ServerName            = $ServerName
+        InstanceName          = $InstanceName
         AvailabilityGroupName = $AvailabilityGroupName
         BackupPath            = $BackupPath
     }

diff --git a/DSCResources/MSFT_SqlAGDatabase/MSFT_SqlAGDatabase.schema.mof b/DSCResources/MSFT_SqlAGDatabase/MSFT_SqlAGDatabase.schema.mof
@@ -9,6 +9,7 @@ class MSFT_SqlAGDatabase : OMI_BaseResource
     [Write, Description("Specifies the membership of the database(s) in the availability group. The options are: Present: The defined database(s) are added to the availability group. All other databases that may be a member of the availability group are ignored. Absent: The defined database(s) are removed from the availability group. All other databases that may be a member of the availability group are ignored. The default is 'Present'."), ValueMap{"Present","Absent"}, Values{"Present","Absent"}] String Ensure;
     [Write, Description("When used with 'Ensure = 'Present'' it ensures the specified database(s) are the only databases that are a member of the specified Availability Group. This parameter is ignored when 'Ensure' is 'Absent'.")] Boolean Force;
     [Write, Description("If set to $true, this ensures the database owner of the database on the primary replica is the owner of the database on all secondary replicas. This requires the database owner is available as a login on all replicas and that the PsDscRunAsCredential has impersonate any login, control server, impersonate login, or control login permissions. If set to $false, the owner of the database will be the PsDscRunAsCredential. The default is '$false'")] Boolean MatchDatabaseOwner;
+    [Write, Description("If set to $true, this adds the restore option WITH REPLACE. If set to $false, Existing databases and files will block the restore and throw error. The default is '$false'.")] Boolean ReplaceExisting;
     [Write, Description("Specifies that the resource will only determine if a change is needed if the target node is the active host of the SQL Server Instance.")] Boolean ProcessOnlyOnActiveNode;
     [Read, Description("Determines if the current node is actively hosting the SQL Server instance.")] Boolean IsActiveNode;
 };
diff --git a/Modules/SqlServerDsc.Common/SqlServerDsc.Common.psm1 b/Modules/SqlServerDsc.Common/SqlServerDsc.Common.psm1
@@ -932,6 +932,9 @@ function Start-SqlSetupProcess
         If the SetupCredential is set, specify with this parameter, which type
         of credentials are set: Native SQL login or Windows user Login. Default
         value is 'WindowsUser'.
+
+    .PARAMETER StatementTimeout
+        Set the query StatementTimeout in seconds. Default 600 seconds (10mins).
 #>
 function Connect-SQL
 {
@@ -956,7 +959,12 @@ function Connect-SQL
         [Parameter()]
         [ValidateSet('WindowsUser', 'SqlLogin')]
         [System.String]
-        $LoginType = 'WindowsUser'
+        $LoginType = 'WindowsUser',
+
+        [Parameter()]
+        [ValidateNotNull()]
+        [System.Int32]
+        $StatementTimeout = 600
     )
 
     Import-SQLPSModule
@@ -970,10 +978,13 @@ function Connect-SQL
         $databaseEngineInstance = "$ServerName\$InstanceName"
     }
 
+    $sql = New-Object -TypeName Microsoft.SqlServer.Management.Smo.Server
+    $sql.ConnectionContext.ServerInstance = $databaseEngineInstance
+    $sql.ConnectionContext.StatementTimeout = $StatementTimeout
+    $sql.ConnectionContext.ApplicationName = 'SqlServerDsc'
+
     if ($SetupCredential)
     {
-        $sql = New-Object -TypeName Microsoft.SqlServer.Management.Smo.Server
-
         if ($LoginType -eq 'SqlLogin')
         {
             $connectUsername = $SetupCredential.Username
@@ -996,21 +1007,23 @@ function Connect-SQL
             'Connecting using the credential ''{0}'' and the login type ''{1}''.' `
                 -f $connectUsername, $LoginType
         ) -Verbose
-
-        $sql.ConnectionContext.ServerInstance = $databaseEngineInstance
-        $sql.ConnectionContext.Connect()
-    }
-    else
-    {
-        $sql = New-Object -TypeName Microsoft.SqlServer.Management.Smo.Server -ArgumentList $databaseEngineInstance
     }
 
-    if ( $sql.Status -match '^Online$' )
+    try
     {
-        Write-Verbose -Message ($script:localizedData.ConnectedToDatabaseEngineInstance -f $databaseEngineInstance) -Verbose
-        return $sql
+        $sql.ConnectionContext.Connect()
+
+        if ( $sql.Status -match '^Online$' )
+        {
+            Write-Verbose -Message ($script:localizedData.ConnectedToDatabaseEngineInstance -f $databaseEngineInstance) -Verbose
+            return $sql
+        }
+        else
+        {
+            throw
+        }
     }
-    else
+    catch
     {
         $errorMessage = $script:localizedData.FailedToConnectToDatabaseEngineInstance -f $databaseEngineInstance
         New-InvalidOperationException -Message $errorMessage
@@ -1538,6 +1551,9 @@ function Restart-ReportingServicesService
     .PARAMETER WithResults
     Specifies if the query should return results.
 
+    .PARAMETER StatementTimeout
+    Set the query StatementTimeout in seconds. Default 600 seconds (10mins).
+
     .EXAMPLE
     Invoke-Query -SQLServer Server1 -SQLInstanceName MSSQLSERVER -Database master -Query 'SELECT name FROM sys.databases' -WithResults
 
@@ -1568,10 +1584,15 @@ function Invoke-Query
 
         [Parameter()]
         [Switch]
-        $WithResults
+        $WithResults,
+
+        [Parameter()]
+        [ValidateNotNull()]
+        [System.Int32]
+        $StatementTimeout = 600
     )
 
-    $serverObject = Connect-SQL -ServerName $SQLServer -InstanceName $SQLInstanceName
+    $serverObject = Connect-SQL -ServerName $SQLServer -InstanceName $SQLInstanceName -StatementTimeout $StatementTimeout
 
     if ( $WithResults )
     {

diff --git a/README.md b/README.md
@@ -294,6 +294,11 @@ group.
   login, control server, impersonate login, or control login permissions.
   If set to $false, the owner of the database will be the PsDscRunAsCredential.
   The default is '$false'.
+* **`[Boolean]` ReplaceExisting** _(Write)_: If set to $true, this adds the restore
+  option WITH REPLACE.
+  If set to $false, Existing databases and files will block the restore
+  and throw error.
+  The default is '$false'.
 * **`[Boolean]` ProcessOnlyOnActiveNode** _(Write)_: Specifies that the resource
   will only determine if a change is needed if the target node is the active
   host of the SQL Server Instance.

diff --git a/Tests/Unit/MSFT_SqlAGDatabase.Tests.ps1 b/Tests/Unit/MSFT_SqlAGDatabase.Tests.ps1
@@ -476,6 +476,7 @@ REVERT'
                     Ensure = 'Present'
                     Force = $false
                     MatchDatabaseOwner = $true
+                    ReplaceExisting = $false
                 }
 
                 Mock -CommandName Add-SqlAvailabilityDatabase -Verifiable -ParameterFilter { $InputObject.PrimaryReplicaServerName -eq 'Server1' -and $InputObject.LocalReplicaRole -eq 'Primary' }
@@ -599,6 +600,34 @@ REVERT'
                     Assert-MockCalled -CommandName Test-ImpersonatePermissions -Scope It -Times 0 -Exactly
                 }
 
+                It 'Should add the specified databases to the availability group when "ReplaceExisting" is $true' {
+                    $mockSetTargetResourceParameters.DatabaseName = 'DB1'
+                    $mockSetTargetResourceParameters.ReplaceExisting = $true
+
+                    { Set-TargetResource @mockSetTargetResourceParameters } | Should -Not -Throw
+
+                    Assert-MockCalled -CommandName Add-SqlAvailabilityDatabase -Scope It -Times 1 -Exactly -ParameterFilter { $InputObject.PrimaryReplicaServerName -eq 'Server1' -and $InputObject.LocalReplicaRole -eq 'Primary' }
+                    Assert-MockCalled -CommandName Add-SqlAvailabilityDatabase -Scope It -Times 1 -Exactly -ParameterFilter { $InputObject.PrimaryReplicaServerName -eq 'Server1' -and $InputObject.LocalReplicaRole -eq 'Secondary' }
+                    Assert-MockCalled -CommandName Add-SqlAvailabilityDatabase -Scope It -Times 0 -Exactly -ParameterFilter { $InputObject.PrimaryReplicaServerName -eq 'Server2' -and $InputObject.LocalReplicaRole -eq 'Primary' }
+                    Assert-MockCalled -CommandName Add-SqlAvailabilityDatabase -Scope It -Times 0 -Exactly -ParameterFilter { $InputObject.PrimaryReplicaServerName -eq 'Server2' -and $InputObject.LocalReplicaRole -eq 'Secondary' }
+                    Assert-MockCalled -CommandName Backup-SqlDatabase -Scope It -Times 1 -Exactly -ParameterFilter { $BackupAction -eq 'Database' }
+                    Assert-MockCalled -CommandName Backup-SqlDatabase -Scope It -Times 1 -Exactly -ParameterFilter { $BackupAction -eq 'Log' }
+                    Assert-MockCalled -CommandName Connect-SQL -Scope It -Times 1 -Exactly -ParameterFilter { $ServerName -eq 'Server1' -and $InstanceName -eq 'MSSQLSERVER' }
+                    Assert-MockCalled -CommandName Connect-SQL -Scope It -Times 1 -Exactly -ParameterFilter { $ServerName -eq 'Server1' }
+                    Assert-MockCalled -CommandName Connect-SQL -Scope It -Times 3 -Exactly -ParameterFilter { $ServerName -eq 'Server2' }
+                    Assert-MockCalled -CommandName Get-PrimaryReplicaServerObject -Scope It -Times 1 -Exactly -ParameterFilter { $AvailabilityGroup.PrimaryReplicaServerName -eq 'Server1' }
+                    Assert-MockCalled -CommandName Get-PrimaryReplicaServerObject -Scope It -Times 0 -Exactly -ParameterFilter { $AvailabilityGroup.PrimaryReplicaServerName -eq 'Server2' }
+                    Assert-MockCalled -CommandName Import-SQLPSModule -Scope It -Times 1 -Exactly
+                    Assert-MockCalled -CommandName Invoke-Query -Scope It -Times 2 -Exactly -ParameterFilter { $Query -like 'EXEC master.dbo.xp_fileexist *' }
+                    Assert-MockCalled -CommandName Invoke-Query -Scope It -Times 0 -Exactly -ParameterFilter $mockInvokeQueryParameterRestoreDatabase
+                    Assert-MockCalled -CommandName Invoke-Query -Scope It -Times 2 -Exactly -ParameterFilter $mockInvokeQueryParameterRestoreDatabaseWithExecuteAs
+                    Assert-MockCalled -CommandName Join-Path -Scope It -Times 1 -Exactly -ParameterFilter { $ChildPath -like '*_Full_*.bak' }
+                    Assert-MockCalled -CommandName Join-Path -Scope It -Times 1 -Exactly -ParameterFilter { $ChildPath -like '*_Log_*.trn' }
+                    Assert-MockCalled -CommandName Remove-Item -Scope It -Times 1 -Exactly
+                    Assert-MockCalled -CommandName Remove-SqlAvailabilityDatabase -Scope It -Times 0 -Exactly
+                    Assert-MockCalled -CommandName Test-ImpersonatePermissions -Scope It -Times 1 -Exactly
+                }
+
                 It 'Should throw the correct error when "MatchDatabaseOwner" is $true and the current login does not have impersonate permissions' {
                     Mock -CommandName Test-ImpersonatePermissions -MockWith { $false } -Verifiable
 

diff --git a/Tests/Unit/SqlServerDsc.Common.Tests.ps1 b/Tests/Unit/SqlServerDsc.Common.Tests.ps1
@@ -2238,6 +2238,8 @@ InModuleScope 'SqlServerDsc.Common' {
                             Add-Member -MemberType NoteProperty -Name ConnectAsUser -Value $false -PassThru |
                             Add-Member -MemberType NoteProperty -Name ConnectAsUserPassword -Value '' -PassThru |
                             Add-Member -MemberType NoteProperty -Name ConnectAsUserName -Value '' -PassThru |
+                            Add-Member -MemberType NoteProperty -Name StatementTimeout -Value 600 -PassThru |
+                            Add-Member -MemberType NoteProperty -Name ApplicationName -Value 'SqlServerDsc' -PassThru |
                             Add-Member -MemberType ScriptMethod -Name Connect -Value {
                                 if ($mockExpectedDatabaseEngineInstance -eq 'MSSQLSERVER')
                                 {
@@ -2386,6 +2388,7 @@ InModuleScope 'SqlServerDsc.Common' {
                 $mockExpectedDatabaseEngineInstance = $mockInstanceName
 
                 Mock -CommandName New-Object `
+                    -MockWith $mockNewObject_MicrosoftDatabaseEngine `
                     -ParameterFilter $mockNewObject_MicrosoftDatabaseEngine_ParameterFilter `
                     -Verifiable