cmbrose · cmbrose · Apr 9, 2024 · Apr 9, 2024 · Apr 9, 2024 · Apr 9, 2024
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,16 @@
 # Release Notes
 All notable changes and release history of the "cosmos-db" module will be documented in this file.
 
+## 1.18
+* Fixes a bug in commands like `Search-CosmosDbRecords` and `Get-AllCosmosDbRecords` which might run for long enough that their auth tokens expire and aren't refreshed. Auth tokens will now be refreshed every 10 min as these commands run.
+* Adds a `-enableAuthHeaderReuse` flag to `Use-CosmosDbInternalFlag` which disables the 10 minute refresh period and forces auth header refreshes for every API call.
+
+## 1.17
+* Fixes `Search-CosmosDbRecords` for partition key range uses where the PK range fetch call didn't use continuation tokens and might miss some results.
+
+## 1.16
+* Adds support for readonly keys via `Use-CosmosDbReadonlyKeys`
+
 ## 1.15
 * Adds support for optimistic concurrency (enabled by default) to `Update-CosmosDbRecord`
 

diff --git a/README.md b/README.md
@@ -263,8 +263,9 @@ Use-CosmosDbInternalFlag -EnableFiddlerDebugging $true
 | Name | Usage | Required |
 | - | - | - |
 | EnableFiddlerDebugging | Sets the `az` flag `env:AZURE_CLI_DISABLE_CONNECTION_VERIFICATION` which enables `az` commands with a Fiddler proxy | No - default is disabled |
-| EnableCaching | Enables caching certain values like DB keys, partition ranges, etc. Improves performance of nearly all operations. | No - default is enabled |
-| EnablePartitionKeyRangeSearches | **[Experimental]** Enables filtering `Search` queries to relevant partition ranges instead of a full scan. Improves performance of `Search` commands. | No - default is disabled |
+| EnableCaching | Enables/disables caching certain values like DB keys, partition ranges, etc. Improves performance of nearly all operations. | No - default is enabled |
+| EnablePartitionKeyRangeSearches | **[Experimental]** Enables/disables filtering `Search` queries to relevant partition ranges instead of a full scan. Improves performance of `Search` commands. | No - default is disabled |
+| EnableAuthHeaderReuse | Enables/disables reusing auth headers for commands which use continuation tokens, like `Search-CosmosDbRecords` or `Get-AllCosmosDbRecords`. | No - default is enabled |
 
 ## Error Handling
 

diff --git a/cosmos-db/cosmos-db.psd1 b/cosmos-db/cosmos-db.psd1
@@ -11,7 +11,7 @@
     # RootModule = ''
 
     # Version number of this module.
-    ModuleVersion     = '1.17'
+    ModuleVersion     = '1.18'
 
     # Supported PSEditions
     # CompatiblePSEditions = @()

diff --git a/cosmos-db/cosmos-db.psm1 b/cosmos-db/cosmos-db.psm1
@@ -10,6 +10,10 @@ $DELETE_VERB = "delete"
 
 $API_VERSION = "2018-12-31"
 
+# Authorization headers are valid for 15 minutes
+# https://learn.microsoft.com/en-us/rest/api/storageservices/authorize-with-shared-key#specifying-the-date-header
+$AUTHORIZATION_HEADER_REFRESH_THRESHOLD = [System.TimeSpan]::FromMinutes(10);
+
 $MASTER_KEY_CACHE = @{}
 $SIGNATURE_HASH_CACHE = @{}
 $PARTITION_KEY_RANGE_CACHE = @{}
@@ -141,13 +145,13 @@ Function Get-AuthorizationHeader([string]$ResourceGroup, [string]$SubscriptionId
 
 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;
-        "Authorization" = $encodedAuthString;
         "Cache-Control" = "No-Cache";
         "Content-Type"  = $contentType;
     }
 
+    Set-AuthHeaders -headers $headers -now $now -encodedAuthString $encodedAuthString
+
     if ($isQuery) {
         $headers["x-ms-documentdb-isquery"] = "true"
     }
@@ -163,6 +167,16 @@ Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$co
     $headers
 }
 
+Function Set-AuthHeaders($headers, [string]$now, [string]$encodedAuthString) {
+    if ($now) {
+        $headers["x-ms-date"] = $now
+    }
+
+    if ($encodedAuthString) {
+        $headers["Authorization"] = $encodedAuthString
+    }
+}
+
 Function Get-QueryParametersAsNameValuePairs($obj) {
     if (!$obj) {
         return @()
@@ -239,17 +253,28 @@ Function Get-ContinuationToken($response) {
     }
 }
 
-Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url, $headers, $body = $null) {
+Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url, $headers, [ScriptBlock]$refreshAuthHeaders, $body = $null) {
     # Remove in case the headers are reused between multiple calls to this function
     $headers.Remove("x-ms-continuation");
 
+    $authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
+    Set-AuthHeaders -headers $headers -now $authHeaders.now -encodedAuthString $authHeaders.encodedAuthString
+
     $response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
     $response
 
     $continuationToken = Get-ContinuationToken $response
     while ($continuationToken) {
         $headers["x-ms-continuation"] = $continuationToken
 
+        $authHeaderReuseDisabled = $env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE -eq 0
+        $authHeaderExpired = [System.DateTime]::Parse($authHeaders.now) + $AUTHORIZATION_HEADER_REFRESH_THRESHOLD -lt [System.DateTime]::UtcNow
+
+        if ($authHeaderReuseDisabled -or $authHeaderExpired) {
+            $authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
+            Set-AuthHeaders -headers $headers -now $authHeaders.now -encodedAuthString $authHeaders.encodedAuthString
+        }
+
         $response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
         $response
 
@@ -278,14 +303,17 @@ Function Get-PartitionKeyRangesOrError
             return $cacheResult
         }
 
-        $now = Get-Time
-
-        $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $PARTITIONKEYRANGE_TYPE -resourceUrl $collectionsUrl -now $now
+        $refreshAuthHeaders = {
+            $now = Get-Time
+            $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $PARTITIONKEYRANGE_TYPE -resourceUrl $collectionsUrl -now $now
 
-        $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey
+            return @{ now = $now; encodedAuthString = $encodedAuthString }
+        }
+
+        $headers = Get-CommonHeaders -PartitionKey $requestPartitionKey
         $headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
 
-        $response = Invoke-CosmosDbApiRequestWithContinuation -Verb $GET_VERB -Url $url -Headers $headers | Get-CosmosDbRecordContent
+        $response = Invoke-CosmosDbApiRequestWithContinuation -Verb $GET_VERB -Url $url -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders | Get-CosmosDbRecordContent
 
         $ranges = $response.partitionKeyRanges
 
@@ -450,16 +478,19 @@ Function Get-AllCosmosDbRecords(
 
     $url = "$baseUrl/$docsUrl"
 
-    $now = Get-Time
+    $refreshAuthHeaders = {
+        $now = Get-Time
+        $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
 
-    $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
+        return @{ now = $now; encodedAuthString = $encodedAuthString }
+    }
 
     $tmp = $ProgressPreference
     $ProgressPreference = 'SilentlyContinue'
     try {
-        $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true
+        $headers = Get-CommonHeaders -isQuery $true
 
-        Invoke-CosmosDbApiRequestWithContinuation -verb $GET_VERB -url $url -Headers $headers
+        Invoke-CosmosDbApiRequestWithContinuation -verb $GET_VERB -url $url -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders
     }
     catch {
         Get-ExceptionResponseOrThrow $_
@@ -542,18 +573,21 @@ Function Search-CosmosDbRecords(
 ) {
     $Parameters = @(Get-QueryParametersAsNameValuePairs $Parameters)
 
+    if (!$DisableExtraFeatures) {
+        return Search-CosmosDbRecordsWithExtraFeatures -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -Query $Query -Parameters $Parameters -SubscriptionId $SubscriptionId
+    }
+
     $baseUrl = Get-BaseDatabaseUrl $Database
     $collectionsUrl = Get-CollectionsUrl $Container $Collection
     $docsUrl = "$collectionsUrl/$DOCS_TYPE"
 
     $url = "$baseUrl/$docsUrl"
 
-    $now = Get-Time
-
-    $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
+    $refreshAuthHeaders = {
+        $now = Get-Time
+        $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
 
-    if (!$DisableExtraFeatures) {
-        return Search-CosmosDbRecordsWithExtraFeatures -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -Query $Query -Parameters $Parameters -SubscriptionId $SubscriptionId
+        return @{ now = $now; encodedAuthString = $encodedAuthString }
     }
 
     try {
@@ -562,10 +596,10 @@ Function Search-CosmosDbRecords(
             parameters = $Parameters;
         }
 
-        $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
+        $headers = Get-CommonHeaders -isQuery $true -contentType "application/Query+json"
         $headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
 
-        Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
+        Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders
     }
     catch {
         Get-ExceptionResponseOrThrow $_
@@ -590,15 +624,18 @@ Function Search-CosmosDbRecordsWithExtraFeatures
 
     $url = "$baseUrl/$docsUrl"
 
-    $now = Get-Time
-
-    $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
-
     $allPartitionKeyRangesOrError = Get-PartitionKeyRangesOrError -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -SubscriptionId $SubscriptionId
 
     if ($allPartitionKeyRangesOrError.ErrorRecord) {
         return Get-ExceptionResponseOrThrow $allPartitionKeyRangesOrError.ErrorRecord
     }
+
+    $refreshAuthHeaders = {
+        $now = Get-Time
+        $encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
+
+        return @{ now = $now; encodedAuthString = $encodedAuthString }
+    }
 
     try {
         $ranges = $allPartitionKeyRangesOrError.Ranges
@@ -608,7 +645,8 @@ Function Search-CosmosDbRecordsWithExtraFeatures
             parameters = $Parameters;
         }
 
-        $headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
+        $authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
+        $headers = Get-CommonHeaders -now $authHeaders.now -encodedAuthString $authHeaders.encodedAuthString -isQuery $true -contentType "application/Query+json"
         $headers += @{
             "x-ms-documentdb-query-enablecrosspartition"           = "true";
             "x-ms-cosmos-supported-query-features"                 = "NonValueAggregate, Aggregate, Distinct, MultipleOrderBy, OffsetAndLimit, OrderBy, Top, CompositeAggregate, GroupBy, MultipleAggregates";
@@ -640,7 +678,7 @@ Function Search-CosmosDbRecordsWithExtraFeatures
         foreach ($partitionKeyRange in $partitionKeyRanges) {
             $headers["x-ms-documentdb-partitionkeyrangeid"] = $partitionKeyRange.id
 
-            Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
+            Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders
         }
     }
     catch {
@@ -966,7 +1004,8 @@ Function Use-CosmosDbInternalFlag
 (
     $enableFiddlerDebugging = $null,
     $enableCaching = $null,
-    $enablePartitionKeyRangeSearches = $null
+    $enablePartitionKeyRangeSearches = $null,
+    $enableAuthHeaderReuse = $null
 ) {
     if ($null -ne $enableFiddlerDebugging) {
         $env:AZURE_CLI_DISABLE_CONNECTION_VERIFICATION = if ($enableFiddlerDebugging) { 1 } else { 0 }
@@ -979,6 +1018,10 @@ Function Use-CosmosDbInternalFlag
     if ($null -ne $enablePartitionKeyRangeSearches) {
         $env:COSMOS_DB_FLAG_ENABLE_PARTITION_KEY_RANGE_SEARCHES = if ($enablePartitionKeyRangeSearches) { 1 } else { 0 }
     }
+
+    if ($null -ne $enableAuthHeaderReuse) {
+        $env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE = if ($enableAuthHeaderReuse) { 1 } else { 0 }
+    }
 }
 
 Function Use-CosmosDbReadonlyKeys

diff --git a/tests/Get-AllCosmosDbRecords.Tests.ps1 b/tests/Get-AllCosmosDbRecords.Tests.ps1
@@ -18,8 +18,7 @@ InModuleScope cosmos-db {
 
             $MOCK_AUTH_HEADER = "MockAuthHeader"
 
-            Function VerifyGetAuthHeader($ResourceGroup, $SubscriptionId, $Database, $verb, $resourceType, $resourceUrl, $now)
-            {
+            Function VerifyGetAuthHeader($ResourceGroup, $SubscriptionId, $Database, $verb, $resourceType, $resourceUrl, $now) {
                 $ResourceGroup | Should -Be $MOCK_RG
                 $SubscriptionId | Should -Be $MOCK_SUB
 
@@ -28,15 +27,19 @@ InModuleScope cosmos-db {
                 $resourceUrl | Should -Be "dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION"
             }
 
-            Function VerifyInvokeCosmosDbApiRequest($verb, $url, $body, $headers, $partitionKey=$MOCK_RECORD_ID)
-            {
+            Function VerifyInvokeCosmosDbApiRequest($verb, $url, $body, $headers, $refreshAuthHeaders, $partitionKey = $MOCK_RECORD_ID) {
                 $verb | Should -Be "get"
                 $url | Should -Be "https://$MOCK_DB.documents.azure.com/dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION/docs"        
                 $body | Should -Be $null
+
+                $authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
 
                 $global:capturedNow | Should -Not -Be $null
 
-                $expectedHeaders = Get-CommonHeaders -now $global:capturedNow -encodedAuthString $MOCK_AUTH_HEADER -isQuery $true
+                $authHeaders.now | Should -Be $global:capturedNow
+                $authHeaders.encodedAuthString | Should -Be $MOCK_AUTH_HEADER
+
+                $expectedHeaders = Get-CommonHeaders -isQuery $true
 
                 AssertHashtablesEqual $expectedHeaders $headers
             }
@@ -55,14 +58,14 @@ InModuleScope cosmos-db {
         It "Sends correct request" {    
             $response = @{
                 StatusCode = 200;
-                Content = "{}";
-                Headers = @{};
+                Content    = "{}";
+                Headers    = @{};
             }
 
             Mock Invoke-CosmosDbApiRequestWithContinuation {
-                param($verb, $url, $body, $headers) 
+                param($verb, $url, $body, $headers, $refreshAuthHeaders) 
 
-                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
+                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers $refreshAuthHeaders | Out-Null
 
                 $response
             }
@@ -77,26 +80,26 @@ InModuleScope cosmos-db {
         It "Returns multiple responses" {
             $response1 = @{
                 StatusCode = 200;
-                Content = "1";
-                Headers = @{};
+                Content    = "1";
+                Headers    = @{};
             }
 
             $response2 = @{
                 StatusCode = 200;
-                Content = "1";
-                Headers = @{};
+                Content    = "1";
+                Headers    = @{};
             }
 
             $response3 = @{
                 StatusCode = 200;
-                Content = "1";
-                Headers = @{};
+                Content    = "1";
+                Headers    = @{};
             }
 
             Mock Invoke-CosmosDbApiRequestWithContinuation {
-                param($verb, $url, $body, $headers) 
+                param($verb, $url, $body, $headers, $refreshAuthHeaders) 
 
-                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
+                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers $refreshAuthHeaders | Out-Null
 
                 $response1
                 $response2
@@ -118,10 +121,10 @@ InModuleScope cosmos-db {
 
             $recordResponse = [PSCustomObject]@{}
 
-            Mock Invoke-CosmosDbApiRequest {
-                param($verb, $url, $body, $headers) 
+            Mock Invoke-CosmosDbApiRequestWithContinuation {
+                param($verb, $url, $body, $headers, $refreshAuthHeaders) 
 
-                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
+                VerifyInvokeCosmosDbApiRequest $verb $url $body $headers $refreshAuthHeaders | Out-Null
 
                 throw [System.Net.WebException]::new("", $null, [System.Net.WebExceptionStatus]::UnknownError, $response)
             }

diff --git a/tests/Get-CommonHeaders.Tests.ps1 b/tests/Get-CommonHeaders.Tests.ps1
@@ -34,5 +34,15 @@ InModuleScope cosmos-db {
                 "x-ms-documentdb-partitionkey" = "[`"MOCK_PARTITION_KEY`"]";
             } $result
         }
+
+        It "Doesn't require authorization headers" {    
+            $result = Get-CommonHeaders
+
+            AssertHashtablesEqual @{
+                "x-ms-version" = "2018-12-31";
+                "Cache-Control" = "No-Cache";
+                "Content-Type" = "application/json";
+            } $result
+        }
     }
 }