Add retry support for document cmdlets

source/Private/utils/Invoke-CosmosDbRequest.ps1

@@ -64,7 +64,15 @@ function Invoke-CosmosDbRequest
         [Parameter()]
         [ValidateSet('Default', 'UTF-8')]
         [System.String]
-        $Encoding = 'Default'
+        $Encoding = 'Default',
+
+        [Parameter()]
+        [System.Int32]
+        $MaximumRetryCount,
+
+        [Parameter()]
+        [System.Int32]
+        $RetryIntervalSec
     )
 
     if ($PSCmdlet.ParameterSetName -eq 'Account')
@@ -258,6 +266,21 @@ function Invoke-CosmosDbRequest
                     continue
                 }
             }
+            elseif (
+                $_.Exception.Response.StatusCode -ge 400 -and
+                $_.Exception.Response.StatusCode -le 599 -and
+                $MaximumRetryCount -gt $retry
+            )
+            {
+                <#
+                    The exception was caused by an error either in the client request or on the server
+                    If additional retries are specified, retry after RetryIntervalSec until Max count is reached
+                #>
+                $retry++
+                Write-Verbose -Message $($LocalizedData.WebRequestErrorDelay -f $MaximumRetryCount, $RetryIntervalSec)
+                Start-Sleep $RetryIntervalSec
+                continue
+            }
 
             if ($_.Exception.Response)
             {

source/Public/documents/Get-CosmosDbDocumentJson.ps1

@@ -89,7 +89,15 @@ function Get-CosmosDbDocumentJson
         [Alias("ResultHeaders")]
         [Parameter()]
         [ref]
-        $ResponseHeader
+        $ResponseHeader,
+
+        [Parameter()]
+        [System.Int32]
+        $MaximumRetryCount,
+
+        [Parameter()]
+        [System.Int32]
+        $RetryIntervalSec = 5
     )
 
     $null = $PSBoundParameters.Remove('Id')
@@ -196,12 +204,16 @@ function Get-CosmosDbDocumentJson
             Because the headers of this request will contain important information
             then we need to use a plain web request.
         #>
-        $result = Invoke-CosmosDbRequest @PSBoundParameters `
-            -Method $method `
-            -ResourceType 'docs' `
-            -ResourcePath $resourcePath `
-            -Headers $headers `
-            -Body $body
+        $invokeCosmosDbRequest_parameters = @{
+            Method            = $method
+            ResourceType      = 'docs'
+            ResourcePath      = $resourcePath
+            Headers           = $headers
+            Body              = $body
+            MaximumRetryCount = $MaximumRetryCount
+            RetryIntervalSec  = $RetryIntervalSec
+        }
+        $result = Invoke-CosmosDbRequest @PSBoundParameters @invokeCosmosDbRequest_parameters
 
         if ($ResponseHeaderPassed)
         {
@@ -220,11 +232,15 @@ function Get-CosmosDbDocumentJson
             $null = $PSBoundParameters.Remove('PartitionKey')
         }
 
-        $result = Invoke-CosmosDbRequest @PSBoundParameters `
-            -Method $method `
-            -Headers $headers `
-            -ResourceType 'docs' `
-            -ResourcePath ('{0}/{1}' -f $resourcePath, $Id)
+        $invokeCosmosDbRequest_parameters = @{
+            Method            = $method
+            Headers           = $headers
+            ResourceType      = 'docs'
+            ResourcePath      = ('{0}/{1}' -f $resourcePath, $Id)
+            MaximumRetryCount = $MaximumRetryCount
+            RetryIntervalSec  = $RetryIntervalSec
+        }
+        $result = Invoke-CosmosDbRequest @PSBoundParameters @invokeCosmosDbRequest_parameters
     }
 
     $documents = Repair-CosmosDbDocumentEncoding -Content $result.Content

source/Public/documents/New-CosmosDbDocument.ps1

@@ -62,7 +62,15 @@ function New-CosmosDbDocument
 
         [Parameter()]
         [switch]
-        $ReturnJson
+        $ReturnJson,
+
+        [Parameter()]
+        [System.Int32]
+        $MaximumRetryCount,
+
+        [Parameter()]
+        [System.Int32]
+        $RetryIntervalSec = 5
     )
 
     $null = $PSBoundParameters.Remove('CollectionId')
@@ -98,12 +106,16 @@ function New-CosmosDbDocument
         $null = $PSBoundParameters.Remove('PartitionKey')
     }
 
-    $result = Invoke-CosmosDbRequest @PSBoundParameters `
-        -Method 'Post' `
-        -ResourceType 'docs' `
-        -ResourcePath $resourcePath `
-        -Body $DocumentBody `
-        -Headers $headers
+    $invokeCosmosDbRequest_parameters = @{
+        Method            = 'Post'
+        ResourceType      = 'docs'
+        ResourcePath      = $resourcePath
+        Body              = $DocumentBody
+        Headers           = $headers
+        MaximumRetryCount = $MaximumRetryCount
+        RetryIntervalSec  = $RetryIntervalSec
+    }
+    $result = Invoke-CosmosDbRequest @PSBoundParameters @invokeCosmosDbRequest_parameters
 
     if ($ReturnJson.IsPresent)
     {

source/Public/documents/Remove-CosmosDbDocument.ps1

@@ -43,7 +43,15 @@ function Remove-CosmosDbDocument
         [Parameter()]
         [ValidateNotNullOrEmpty()]
         [System.Object[]]
-        $PartitionKey
+        $PartitionKey,
+
+        [Parameter()]
+        [System.Int32]
+        $MaximumRetryCount,
+
+        [Parameter()]
+        [System.Int32]
+        $RetryIntervalSec = 5
     )
 
     $null = $PSBoundParameters.Remove('CollectionId')
@@ -60,10 +68,13 @@ function Remove-CosmosDbDocument
         }
         $null = $PSBoundParameters.Remove('PartitionKey')
     }
-
-    $null = Invoke-CosmosDbRequest @PSBoundParameters `
-        -Method 'Delete' `
-        -ResourceType 'docs' `
-        -ResourcePath $resourcePath `
-        -Headers $headers
+    $invokeCosmosDbRequest_parameters = @{
+        Method            = 'Delete'
+        ResourceType      = 'docs'
+        ResourcePath      = $resourcePath
+        Headers           = $headers
+        MaximumRetryCount = $MaximumRetryCount
+        RetryIntervalSec  = $RetryIntervalSec
+    }
+    $null = Invoke-CosmosDbRequest @PSBoundParameters @invokeCosmosDbRequest_parameters
 }

source/Public/documents/Set-CosmosDbDocument.ps1

@@ -69,7 +69,15 @@ function Set-CosmosDbDocument
 
         [Parameter()]
         [switch]
-        $ReturnJson
+        $ReturnJson,
+
+        [Parameter()]
+        [System.Int32]
+        $MaximumRetryCount,
+
+        [Parameter()]
+        [System.Int32]
+        $RetryIntervalSec = 5
     )
 
     $null = $PSBoundParameters.Remove('CollectionId')
@@ -105,12 +113,16 @@ function Set-CosmosDbDocument
         $null = $PSBoundParameters.Remove('ETag')
     }
 
-    $result = Invoke-CosmosDbRequest @PSBoundParameters `
-        -Method 'Put' `
-        -ResourceType 'docs' `
-        -ResourcePath $resourcePath `
-        -Body $DocumentBody `
-        -Headers $headers
+    $invokeCosmosDbRequest_parameters = @{
+        Method            = 'Put'
+        ResourceType      = 'docs'
+        ResourcePath      = $resourcePath
+        Body              = $DocumentBody
+        Headers           = $headers
+        MaximumRetryCount = $MaximumRetryCount
+        RetryIntervalSec  = $RetryIntervalSec
+    }
+    $result = Invoke-CosmosDbRequest @PSBoundParameters @invokeCosmosDbRequest_parameters
 
     if ($ReturnJson.IsPresent)
     {

source/en-US/CosmosDB.strings.psd1

@@ -26,6 +26,7 @@ ConvertFrom-StringData -StringData @'
     BackOffPolicyAppliedPolicyDelay = The {0} back-off policy delay {1}ms will be used because it is longer than the requested delay {2}ms.
     BackOffPolicyAppliedRequestedDelay = The requested delay {2}ms will be used because it is longer than the {0} back-off policy delay {1}ms.
     WaitingBackoffPolicyDelay = The collection has exceeded the provisioned throughput limit but retry {0} will be attempted in {1}ms.
+    WebRequestErrorDelay = The request failed but will be retried up to {0} times in {1}s.
     ErrorAuthorizationKeyEmpty = The authorization key is empty. It must be passed in the context or a valid token context for the resource being accessed must be supplied.
     WarningNewCollectionOfferTypeDeprecated = The 'OfferType' parameter is a legacy parameter and is only supported for backwards compatibility and may be removed in future. It is recommended to use 'OfferThroughput' or 'AutopilotThroughput' instead.
     ErrorNewCollectionOfferParameterConflict = Only one of 'OfferType', OfferThroughput' or 'AutoscaleThroughput' should be specified when creating a new collection.