Skip to content

Commit 2e153a9

Browse files
cmbrosecmbrose
authored
Fix auth tokens expiring for long running commands (#18)
1 parent eeeae6f commit 2e153a9

9 files changed

+473
-108
lines changed

Diff for: ‎CHANGELOG.md

+10
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,16 @@
11
# Release Notes
22
All notable changes and release history of the "cosmos-db" module will be documented in this file.
33

4+
## 1.18
5+
* 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.
6+
* Adds a `-enableAuthHeaderReuse` flag to `Use-CosmosDbInternalFlag` which disables the 10 minute refresh period and forces auth header refreshes for every API call.
7+
8+
## 1.17
9+
* Fixes `Search-CosmosDbRecords` for partition key range uses where the PK range fetch call didn't use continuation tokens and might miss some results.
10+
11+
## 1.16
12+
* Adds support for readonly keys via `Use-CosmosDbReadonlyKeys`
13+
414
## 1.15
515
* Adds support for optimistic concurrency (enabled by default) to `Update-CosmosDbRecord`
616

Diff for: ‎README.md

+3-2
Original file line numberDiff line numberDiff line change
@@ -263,8 +263,9 @@ Use-CosmosDbInternalFlag -EnableFiddlerDebugging $true
263263
| Name | Usage | Required |
264264
| - | - | - |
265265
| EnableFiddlerDebugging | Sets the `az` flag `env:AZURE_CLI_DISABLE_CONNECTION_VERIFICATION` which enables `az` commands with a Fiddler proxy | No - default is disabled |
266-
| EnableCaching | Enables caching certain values like DB keys, partition ranges, etc. Improves performance of nearly all operations. | No - default is enabled |
267-
| EnablePartitionKeyRangeSearches | **[Experimental]** Enables filtering `Search` queries to relevant partition ranges instead of a full scan. Improves performance of `Search` commands. | No - default is disabled |
266+
| EnableCaching | Enables/disables caching certain values like DB keys, partition ranges, etc. Improves performance of nearly all operations. | No - default is enabled |
267+
| 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 |
268+
| EnableAuthHeaderReuse | Enables/disables reusing auth headers for commands which use continuation tokens, like `Search-CosmosDbRecords` or `Get-AllCosmosDbRecords`. | No - default is enabled |
268269

269270
## Error Handling
270271

Diff for: ‎cosmos-db/cosmos-db.psd1

+1-1
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
# RootModule = ''
1212

1313
# Version number of this module.
14-
ModuleVersion = '1.17'
14+
ModuleVersion = '1.18'
1515

1616
# Supported PSEditions
1717
# CompatiblePSEditions = @()

Diff for: ‎cosmos-db/cosmos-db.psm1

+69-26
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,10 @@ $DELETE_VERB = "delete"
1010

1111
$API_VERSION = "2018-12-31"
1212

13+
# Authorization headers are valid for 15 minutes
14+
# https://learn.microsoft.com/en-us/rest/api/storageservices/authorize-with-shared-key#specifying-the-date-header
15+
$AUTHORIZATION_HEADER_REFRESH_THRESHOLD = [System.TimeSpan]::FromMinutes(10);
16+
1317
$MASTER_KEY_CACHE = @{}
1418
$SIGNATURE_HASH_CACHE = @{}
1519
$PARTITION_KEY_RANGE_CACHE = @{}
@@ -141,13 +145,13 @@ Function Get-AuthorizationHeader([string]$ResourceGroup, [string]$SubscriptionId
141145

142146
Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$contentType = "application/json", [bool]$isQuery = $false, [string]$PartitionKey = $null, [string]$Etag = $null) {
143147
$headers = @{
144-
"x-ms-date" = $now;
145148
"x-ms-version" = $API_VERSION;
146-
"Authorization" = $encodedAuthString;
147149
"Cache-Control" = "No-Cache";
148150
"Content-Type" = $contentType;
149151
}
150152

153+
Set-AuthHeaders -headers $headers -now $now -encodedAuthString $encodedAuthString
154+
151155
if ($isQuery) {
152156
$headers["x-ms-documentdb-isquery"] = "true"
153157
}
@@ -163,6 +167,16 @@ Function Get-CommonHeaders([string]$now, [string]$encodedAuthString, [string]$co
163167
$headers
164168
}
165169

170+
Function Set-AuthHeaders($headers, [string]$now, [string]$encodedAuthString) {
171+
if ($now) {
172+
$headers["x-ms-date"] = $now
173+
}
174+
175+
if ($encodedAuthString) {
176+
$headers["Authorization"] = $encodedAuthString
177+
}
178+
}
179+
166180
Function Get-QueryParametersAsNameValuePairs($obj) {
167181
if (!$obj) {
168182
return @()
@@ -239,17 +253,28 @@ Function Get-ContinuationToken($response) {
239253
}
240254
}
241255

242-
Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url, $headers, $body = $null) {
256+
Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url, $headers, [ScriptBlock]$refreshAuthHeaders, $body = $null) {
243257
# Remove in case the headers are reused between multiple calls to this function
244258
$headers.Remove("x-ms-continuation");
245259

260+
$authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
261+
Set-AuthHeaders -headers $headers -now $authHeaders.now -encodedAuthString $authHeaders.encodedAuthString
262+
246263
$response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
247264
$response
248265

249266
$continuationToken = Get-ContinuationToken $response
250267
while ($continuationToken) {
251268
$headers["x-ms-continuation"] = $continuationToken
252269

270+
$authHeaderReuseDisabled = $env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE -eq 0
271+
$authHeaderExpired = [System.DateTime]::Parse($authHeaders.now) + $AUTHORIZATION_HEADER_REFRESH_THRESHOLD -lt [System.DateTime]::UtcNow
272+
273+
if ($authHeaderReuseDisabled -or $authHeaderExpired) {
274+
$authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
275+
Set-AuthHeaders -headers $headers -now $authHeaders.now -encodedAuthString $authHeaders.encodedAuthString
276+
}
277+
253278
$response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
254279
$response
255280

@@ -278,14 +303,17 @@ Function Get-PartitionKeyRangesOrError
278303
return $cacheResult
279304
}
280305

281-
$now = Get-Time
282-
283-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $PARTITIONKEYRANGE_TYPE -resourceUrl $collectionsUrl -now $now
306+
$refreshAuthHeaders = {
307+
$now = Get-Time
308+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $PARTITIONKEYRANGE_TYPE -resourceUrl $collectionsUrl -now $now
284309

285-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey
310+
return @{ now = $now; encodedAuthString = $encodedAuthString }
311+
}
312+
313+
$headers = Get-CommonHeaders -PartitionKey $requestPartitionKey
286314
$headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
287315

288-
$response = Invoke-CosmosDbApiRequestWithContinuation -Verb $GET_VERB -Url $url -Headers $headers | Get-CosmosDbRecordContent
316+
$response = Invoke-CosmosDbApiRequestWithContinuation -Verb $GET_VERB -Url $url -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders | Get-CosmosDbRecordContent
289317

290318
$ranges = $response.partitionKeyRanges
291319

@@ -450,16 +478,19 @@ Function Get-AllCosmosDbRecords(
450478

451479
$url = "$baseUrl/$docsUrl"
452480

453-
$now = Get-Time
481+
$refreshAuthHeaders = {
482+
$now = Get-Time
483+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
454484

455-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
485+
return @{ now = $now; encodedAuthString = $encodedAuthString }
486+
}
456487

457488
$tmp = $ProgressPreference
458489
$ProgressPreference = 'SilentlyContinue'
459490
try {
460-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true
491+
$headers = Get-CommonHeaders -isQuery $true
461492

462-
Invoke-CosmosDbApiRequestWithContinuation -verb $GET_VERB -url $url -Headers $headers
493+
Invoke-CosmosDbApiRequestWithContinuation -verb $GET_VERB -url $url -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders
463494
}
464495
catch {
465496
Get-ExceptionResponseOrThrow $_
@@ -542,18 +573,21 @@ Function Search-CosmosDbRecords(
542573
) {
543574
$Parameters = @(Get-QueryParametersAsNameValuePairs $Parameters)
544575

576+
if (!$DisableExtraFeatures) {
577+
return Search-CosmosDbRecordsWithExtraFeatures -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -Query $Query -Parameters $Parameters -SubscriptionId $SubscriptionId
578+
}
579+
545580
$baseUrl = Get-BaseDatabaseUrl $Database
546581
$collectionsUrl = Get-CollectionsUrl $Container $Collection
547582
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
548583

549584
$url = "$baseUrl/$docsUrl"
550585

551-
$now = Get-Time
552-
553-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
586+
$refreshAuthHeaders = {
587+
$now = Get-Time
588+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
554589

555-
if (!$DisableExtraFeatures) {
556-
return Search-CosmosDbRecordsWithExtraFeatures -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -Query $Query -Parameters $Parameters -SubscriptionId $SubscriptionId
590+
return @{ now = $now; encodedAuthString = $encodedAuthString }
557591
}
558592

559593
try {
@@ -562,10 +596,10 @@ Function Search-CosmosDbRecords(
562596
parameters = $Parameters;
563597
}
564598

565-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
599+
$headers = Get-CommonHeaders -isQuery $true -contentType "application/Query+json"
566600
$headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
567601

568-
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
602+
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders
569603
}
570604
catch {
571605
Get-ExceptionResponseOrThrow $_
@@ -590,15 +624,18 @@ Function Search-CosmosDbRecordsWithExtraFeatures
590624

591625
$url = "$baseUrl/$docsUrl"
592626

593-
$now = Get-Time
594-
595-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
596-
597627
$allPartitionKeyRangesOrError = Get-PartitionKeyRangesOrError -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -SubscriptionId $SubscriptionId
598628

599629
if ($allPartitionKeyRangesOrError.ErrorRecord) {
600630
return Get-ExceptionResponseOrThrow $allPartitionKeyRangesOrError.ErrorRecord
601631
}
632+
633+
$refreshAuthHeaders = {
634+
$now = Get-Time
635+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
636+
637+
return @{ now = $now; encodedAuthString = $encodedAuthString }
638+
}
602639

603640
try {
604641
$ranges = $allPartitionKeyRangesOrError.Ranges
@@ -608,7 +645,8 @@ Function Search-CosmosDbRecordsWithExtraFeatures
608645
parameters = $Parameters;
609646
}
610647

611-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
648+
$authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
649+
$headers = Get-CommonHeaders -now $authHeaders.now -encodedAuthString $authHeaders.encodedAuthString -isQuery $true -contentType "application/Query+json"
612650
$headers += @{
613651
"x-ms-documentdb-query-enablecrosspartition" = "true";
614652
"x-ms-cosmos-supported-query-features" = "NonValueAggregate, Aggregate, Distinct, MultipleOrderBy, OffsetAndLimit, OrderBy, Top, CompositeAggregate, GroupBy, MultipleAggregates";
@@ -640,7 +678,7 @@ Function Search-CosmosDbRecordsWithExtraFeatures
640678
foreach ($partitionKeyRange in $partitionKeyRanges) {
641679
$headers["x-ms-documentdb-partitionkeyrangeid"] = $partitionKeyRange.id
642680

643-
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
681+
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers -RefreshAuthHeaders $refreshAuthHeaders
644682
}
645683
}
646684
catch {
@@ -966,7 +1004,8 @@ Function Use-CosmosDbInternalFlag
9661004
(
9671005
$enableFiddlerDebugging = $null,
9681006
$enableCaching = $null,
969-
$enablePartitionKeyRangeSearches = $null
1007+
$enablePartitionKeyRangeSearches = $null,
1008+
$enableAuthHeaderReuse = $null
9701009
) {
9711010
if ($null -ne $enableFiddlerDebugging) {
9721011
$env:AZURE_CLI_DISABLE_CONNECTION_VERIFICATION = if ($enableFiddlerDebugging) { 1 } else { 0 }
@@ -979,6 +1018,10 @@ Function Use-CosmosDbInternalFlag
9791018
if ($null -ne $enablePartitionKeyRangeSearches) {
9801019
$env:COSMOS_DB_FLAG_ENABLE_PARTITION_KEY_RANGE_SEARCHES = if ($enablePartitionKeyRangeSearches) { 1 } else { 0 }
9811020
}
1021+
1022+
if ($null -ne $enableAuthHeaderReuse) {
1023+
$env:COSMOS_DB_FLAG_ENABLE_AUTH_HEADER_REUSE = if ($enableAuthHeaderReuse) { 1 } else { 0 }
1024+
}
9821025
}
9831026

9841027
Function Use-CosmosDbReadonlyKeys

Diff for: ‎tests/Get-AllCosmosDbRecords.Tests.ps1

+23-20
Original file line numberDiff line numberDiff line change
@@ -18,8 +18,7 @@ InModuleScope cosmos-db {
1818

1919
$MOCK_AUTH_HEADER = "MockAuthHeader"
2020

21-
Function VerifyGetAuthHeader($ResourceGroup, $SubscriptionId, $Database, $verb, $resourceType, $resourceUrl, $now)
22-
{
21+
Function VerifyGetAuthHeader($ResourceGroup, $SubscriptionId, $Database, $verb, $resourceType, $resourceUrl, $now) {
2322
$ResourceGroup | Should -Be $MOCK_RG
2423
$SubscriptionId | Should -Be $MOCK_SUB
2524

@@ -28,15 +27,19 @@ InModuleScope cosmos-db {
2827
$resourceUrl | Should -Be "dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION"
2928
}
3029

31-
Function VerifyInvokeCosmosDbApiRequest($verb, $url, $body, $headers, $partitionKey=$MOCK_RECORD_ID)
32-
{
30+
Function VerifyInvokeCosmosDbApiRequest($verb, $url, $body, $headers, $refreshAuthHeaders, $partitionKey = $MOCK_RECORD_ID) {
3331
$verb | Should -Be "get"
3432
$url | Should -Be "https://$MOCK_DB.documents.azure.com/dbs/$MOCK_CONTAINER/colls/$MOCK_COLLECTION/docs"
3533
$body | Should -Be $null
34+
35+
$authHeaders = Invoke-Command -ScriptBlock $refreshAuthHeaders
3636

3737
$global:capturedNow | Should -Not -Be $null
3838

39-
$expectedHeaders = Get-CommonHeaders -now $global:capturedNow -encodedAuthString $MOCK_AUTH_HEADER -isQuery $true
39+
$authHeaders.now | Should -Be $global:capturedNow
40+
$authHeaders.encodedAuthString | Should -Be $MOCK_AUTH_HEADER
41+
42+
$expectedHeaders = Get-CommonHeaders -isQuery $true
4043

4144
AssertHashtablesEqual $expectedHeaders $headers
4245
}
@@ -55,14 +58,14 @@ InModuleScope cosmos-db {
5558
It "Sends correct request" {
5659
$response = @{
5760
StatusCode = 200;
58-
Content = "{}";
59-
Headers = @{};
61+
Content = "{}";
62+
Headers = @{};
6063
}
6164

6265
Mock Invoke-CosmosDbApiRequestWithContinuation {
63-
param($verb, $url, $body, $headers)
66+
param($verb, $url, $body, $headers, $refreshAuthHeaders)
6467

65-
VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
68+
VerifyInvokeCosmosDbApiRequest $verb $url $body $headers $refreshAuthHeaders | Out-Null
6669

6770
$response
6871
}
@@ -77,26 +80,26 @@ InModuleScope cosmos-db {
7780
It "Returns multiple responses" {
7881
$response1 = @{
7982
StatusCode = 200;
80-
Content = "1";
81-
Headers = @{};
83+
Content = "1";
84+
Headers = @{};
8285
}
8386

8487
$response2 = @{
8588
StatusCode = 200;
86-
Content = "1";
87-
Headers = @{};
89+
Content = "1";
90+
Headers = @{};
8891
}
8992

9093
$response3 = @{
9194
StatusCode = 200;
92-
Content = "1";
93-
Headers = @{};
95+
Content = "1";
96+
Headers = @{};
9497
}
9598

9699
Mock Invoke-CosmosDbApiRequestWithContinuation {
97-
param($verb, $url, $body, $headers)
100+
param($verb, $url, $body, $headers, $refreshAuthHeaders)
98101

99-
VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
102+
VerifyInvokeCosmosDbApiRequest $verb $url $body $headers $refreshAuthHeaders | Out-Null
100103

101104
$response1
102105
$response2
@@ -118,10 +121,10 @@ InModuleScope cosmos-db {
118121

119122
$recordResponse = [PSCustomObject]@{}
120123

121-
Mock Invoke-CosmosDbApiRequest {
122-
param($verb, $url, $body, $headers)
124+
Mock Invoke-CosmosDbApiRequestWithContinuation {
125+
param($verb, $url, $body, $headers, $refreshAuthHeaders)
123126

124-
VerifyInvokeCosmosDbApiRequest $verb $url $body $headers | Out-Null
127+
VerifyInvokeCosmosDbApiRequest $verb $url $body $headers $refreshAuthHeaders | Out-Null
125128

126129
throw [System.Net.WebException]::new("", $null, [System.Net.WebExceptionStatus]::UnknownError, $response)
127130
}

Diff for: ‎tests/Get-CommonHeaders.Tests.ps1

+10
Original file line numberDiff line numberDiff line change
@@ -34,5 +34,15 @@ InModuleScope cosmos-db {
3434
"x-ms-documentdb-partitionkey" = "[`"MOCK_PARTITION_KEY`"]";
3535
} $result
3636
}
37+
38+
It "Doesn't require authorization headers" {
39+
$result = Get-CommonHeaders
40+
41+
AssertHashtablesEqual @{
42+
"x-ms-version" = "2018-12-31";
43+
"Cache-Control" = "No-Cache";
44+
"Content-Type" = "application/json";
45+
} $result
46+
}
3747
}
3848
}

0 commit comments

Comments
 (0)