Add etag verfication in update calls (#14)

aditya30394 · web-flow · commit a07987141d94 · 2023-06-01T20:07:53.000-05:00
diff --git a/cosmos-db/cosmos-db.psm1 b/cosmos-db/cosmos-db.psm1
@@ -123,7 +123,7 @@ Function Get-AuthorizationHeader([string]$ResourceGroup, [string]$SubscriptionId
     Get-EncodedAuthString -signatureHash $signatureHash
 }
 
-Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$contentType = "application/json", [bool]$isQuery = $false, [string]$PartitionKey = $null) {
+Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$contentType = "application/json", [bool]$isQuery = $false, [string]$PartitionKey = $null, [string]$Etag = $null) {
     $headers = @{ 
         "x-ms-date"     = $now;
         "x-ms-version"  = $API_VERSION;
@@ -139,6 +139,10 @@ Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$co
     if ($PartitionKey) {
         $headers["x-ms-documentdb-partitionkey"] = "[`"$PartitionKey`"]"
     }
+	
+	if ($Etag) {
+		$headers["If-Match"] = $Etag
+	}
 
     $headers
 }
@@ -738,6 +742,8 @@ Function New-CosmosDbRecord {
     [Optional] The record's partition key. Default is the `id` property of `Object`. Required if using a custom partition strategy.
 .PARAMETER GetPartitionKeyBlock
     [Optional] Callback to get the partition key from the input object. Default is the `id` property of `Object`. Required if using a custom partition strategy.
+.PARAMETER EnforceOptimisticConcurrency
+    [Optional] Boolean specifying whether to enforce optimistic concurrency. Default is $true.
 
 .EXAMPLE
     $> Update-CosmosDbRecord -Object @{ id = 1234; key = value } ...
@@ -774,7 +780,9 @@ Function Update-CosmosDbRecord {
         [parameter(Mandatory = $true)][string]$Collection, 
         [parameter(Mandatory = $false)][string]$SubscriptionId = "", 
         [parameter(Mandatory = $false, ParameterSetName = "ExplicitPartitionKey")][string]$PartitionKey = "", 
-        [parameter(Mandatory = $false, ParameterSetName = "ParttionKeyCallback")]$GetPartitionKeyBlock = $null
+        [parameter(Mandatory = $false, ParameterSetName = "ParttionKeyCallback")]$GetPartitionKeyBlock = $null,
+        [parameter(Mandatory = $false)][bool]$EnforceOptimisticConcurrency = $true
+		
     )
 
     begin {
@@ -792,7 +800,11 @@ Function Update-CosmosDbRecord {
             
             $requestPartitionKey = if ($PartitionKey) { $PartitionKey } elseif ($GetPartitionKeyBlock) { Invoke-Command -ScriptBlock $GetPartitionKeyBlock -ArgumentList $Object } else { $Object.Id }
 
-            $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey
+            if ($EnforceOptimisticConcurrency) {
+                $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey -Etag $Object._etag
+            } else {
+                $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey
+            }
 
             Invoke-CosmosDbApiRequest -Verb $PUT_VERB -Url $url -Body $Object -Headers $headers
         }
diff --git a/tests/Update-CosmosDbRecord.Tests.ps1 b/tests/Update-CosmosDbRecord.Tests.ps1
@@ -16,6 +16,7 @@ InModuleScope cosmos-db {
             $MOCK_CONTAINER = "MOCK_CONTAINER"
             $MOCK_COLLECTION = "MOCK_COLLECTION"
             $MOCK_RECORD_ID = "MOCK_RECORD_ID"
+            $MOCK_ETAG = "MOCK_ETAG"
 
             $MOCK_AUTH_HEADER = "MockAuthHeader"
 
@@ -29,7 +30,7 @@ InModuleScope cosmos-db {
                 $resourceUrl | Should -Be "dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION/docs/$expectedId"
             }
 
-            Function VerifyInvokeCosmosDbApiRequest($verb, $url, $actualBody, $expectedBody, $headers, $expectedId=$MOCK_RECORD_ID, $expectedPartitionKey=$null)
+            Function VerifyInvokeCosmosDbApiRequest($verb, $url, $actualBody, $expectedBody, $headers, $expectedId=$MOCK_RECORD_ID, $expectedPartitionKey=$null, $enforceOptimisticConcurrency=$true)
             {
                 $verb | Should -Be "put"
                 $url | Should -Be "https://$MOCK_DB.documents.azure.com/dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION/docs/$expectedId"        
@@ -39,8 +40,12 @@ InModuleScope cosmos-db {
                 $global:capturedNow | Should -Not -Be $null
 
                 $expectedPartitionKey = if ($expectedPartitionKey) { $expectedPartitionKey } else { $expectedId }
-                $expectedHeaders = Get-CommonHeaders -now $global:capturedNow -encodedAuthString $MOCK_AUTH_HEADER -PartitionKey $expectedPartitionKey
-            
+                
+                if ($EnforceOptimisticConcurrency) {
+                    $expectedHeaders = $expectedHeaders = Get-CommonHeaders -now $global:capturedNow -encodedAuthString $MOCK_AUTH_HEADER -PartitionKey $expectedPartitionKey -Etag $MOCK_ETAG
+                } else {
+                    $expectedHeaders = Get-CommonHeaders -now $global:capturedNow -encodedAuthString $MOCK_AUTH_HEADER -PartitionKey $expectedPartitionKey
+                }
                 AssertHashtablesEqual $expectedHeaders $headers
             }
 
@@ -65,6 +70,7 @@ InModuleScope cosmos-db {
                 id = $MOCK_RECORD_ID;
                 key1 = "value1";
                 key2 = 2;
+                "_etag" = $MOCK_ETAG;
             }
 
             Mock Invoke-CosmosDbApiRequest {
@@ -92,6 +98,7 @@ InModuleScope cosmos-db {
                 id = $MOCK_RECORD_ID;
                 key1 = "value1";
                 key2 = 2;
+                "_etag" = $MOCK_ETAG;
             }
 
             $MOCK_PARTITION_KEY = "MOCK_PARTITION_KEY"
@@ -121,6 +128,7 @@ InModuleScope cosmos-db {
                 id = $MOCK_RECORD_ID;
                 key1 = "value1";
                 key2 = 2;
+                "_etag" = $MOCK_ETAG;
             }
 
             $MOCK_PARTITION_KEY = "MOCK_PARTITION_KEY"
@@ -147,11 +155,48 @@ InModuleScope cosmos-db {
             Assert-MockCalled Invoke-CosmosDbApiRequest -Times 1
         }
 
+        It "Optimistic concurrency can be disabled" {
+            $response = @{
+                StatusCode = 200;
+                Content = "{}"
+            }
+
+            $payload = @{
+                id = $MOCK_RECORD_ID;
+                key1 = "value1";
+                key2 = 2;
+                "_etag" = $MOCK_ETAG;
+            }
+
+            $MOCK_PARTITION_KEY = "MOCK_PARTITION_KEY"
+            $MOCK_GET_PARTITION_KEY = { 
+                param($obj)
+
+                $obj | Should -BeExactly $payload | Out-Null
+
+                $MOCK_PARTITION_KEY
+            }
+
+            Mock Invoke-CosmosDbApiRequest {
+                param($verb, $url, $body, $headers) 
+                
+                VerifyInvokeCosmosDbApiRequest $verb $url $body $payload $headers -ExpectedPartitionKey $MOCK_PARTITION_KEY -EnforceOptimisticConcurrency $false | Out-Null
+        
+                $response
+            }
+
+            $result = $payload | Update-CosmosDbRecord -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Container $MOCK_CONTAINER -Collection $MOCK_COLLECTION -GetPartitionKeyBlock $MOCK_GET_PARTITION_KEY -EnforceOptimisticConcurrency $false
+
+            $result | Should -BeExactly $response
+
+            Assert-MockCalled Invoke-CosmosDbApiRequest -Times 1
+        }
+
         It "Sends correct request with custom partition key callback for multiple inputs" {
             $payloads = @(
-                @{ id = "1" };
-                @{ id = "2" };
-                @{ id = "3" };
+                @{ id = "1"; "_etag" = $MOCK_ETAG };
+                @{ id = "2"; "_etag" = $MOCK_ETAG };
+                @{ id = "3"; "_etag" = $MOCK_ETAG };
             )
 
             $global:idx = 0
@@ -199,6 +244,7 @@ InModuleScope cosmos-db {
             $result.Count | Should -Be $payloads.Count
 
             Assert-MockCalled Invoke-CosmosDbApiRequest -Times $payloads.Count
+            Assert-MockCalled Get-AuthorizationHeader -Times $payloads.Count
         }
         
         It "Url encodes the record id in the API url" {    
@@ -212,7 +258,8 @@ InModuleScope cosmos-db {
             $expectedAuthHeaderRecordId = $testRecordId # The id in the auth header should not be encoded
 
             $payload = @{
-                id = $testRecordId
+                id = $testRecordId;
+                "_etag" = $MOCK_ETAG;
             }
 
             Mock Invoke-CosmosDbApiRequest {
@@ -236,6 +283,9 @@ InModuleScope cosmos-db {
             $result = $payload | Update-CosmosDbRecord -ResourceGroup $MOCK_RG -SubscriptionId $MOCK_SUB -Database $MOCK_DB -Container $MOCK_CONTAINER -Collection $MOCK_COLLECTION
 
             $result | Should -BeExactly $response
+
+            Assert-MockCalled Invoke-CosmosDbApiRequest -Times 1
+            Assert-MockCalled Get-AuthorizationHeader -Times 1
         }
 
         It "Should handle exceptions gracefully" {    
@@ -247,6 +297,7 @@ InModuleScope cosmos-db {
                 id = $MOCK_RECORD_ID;
                 key1 = "value1";
                 key2 = 2;
+                "_etag" = $MOCK_ETAG;
             }
 
             Mock Invoke-CosmosDbApiRequest {
@@ -269,6 +320,7 @@ InModuleScope cosmos-db {
 
             $result | Should -BeExactly $recordResponse
             Assert-MockCalled Get-ExceptionResponseOrThrow -Times 1
+            Assert-MockCalled Invoke-CosmosDbApiRequest -Times 1
         }
     }
 }
