Skip to content

Commit eeeae6f

Browse files
cmbrosecmbrose
committed
Fix partition key ranges not using continuations, and some pipeline cleanup
1 parent 8389e3c commit eeeae6f

File tree

5 files changed

+279
-169
lines changed

5 files changed

+279
-169
lines changed

‎cosmos-db/cosmos-db.psd1‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
# RootModule = ''
1212

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

1616
# Supported PSEditions
1717
# CompatiblePSEditions = @()

‎cosmos-db/cosmos-db.psm1‎

Lines changed: 123 additions & 130 deletions
Original file line numberDiff line numberDiff line change
@@ -240,19 +240,20 @@ Function Get-ContinuationToken($response) {
240240
}
241241

242242
Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url, $headers, $body = $null) {
243-
process {
243+
# Remove in case the headers are reused between multiple calls to this function
244+
$headers.Remove("x-ms-continuation");
245+
246+
$response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
247+
$response
248+
249+
$continuationToken = Get-ContinuationToken $response
250+
while ($continuationToken) {
251+
$headers["x-ms-continuation"] = $continuationToken
252+
244253
$response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
245254
$response
246255

247256
$continuationToken = Get-ContinuationToken $response
248-
while ($continuationToken) {
249-
$headers["x-ms-continuation"] = $continuationToken
250-
251-
$response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
252-
$response
253-
254-
$continuationToken = Get-ContinuationToken $response
255-
}
256257
}
257258
}
258259

@@ -284,7 +285,7 @@ Function Get-PartitionKeyRangesOrError
284285
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey
285286
$headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
286287

287-
$response = Invoke-CosmosDbApiRequest -Verb $GET_VERB -Url $url -Headers $headers | Get-CosmosDbRecordContent
288+
$response = Invoke-CosmosDbApiRequestWithContinuation -Verb $GET_VERB -Url $url -Headers $headers | Get-CosmosDbRecordContent
288289

289290
$ranges = $response.partitionKeyRanges
290291

@@ -368,29 +369,27 @@ Function Get-CosmosDbRecord(
368369
[parameter(Mandatory = $true)][string]$Collection,
369370
[parameter(Mandatory = $true)][string]$RecordId,
370371
[parameter(Mandatory = $false)][string]$SubscriptionId = "",
371-
[parameter(Mandatory = $false)][string]$PartitionKey = "") {
372-
begin {
373-
$baseUrl = Get-BaseDatabaseUrl $Database
374-
$documentUrls = Get-DocumentsUrl $Container $Collection $RecordId
372+
[parameter(Mandatory = $false)][string]$PartitionKey = ""
373+
) {
374+
$baseUrl = Get-BaseDatabaseUrl $Database
375+
$documentUrls = Get-DocumentsUrl $Container $Collection $RecordId
375376

376-
$url = "$baseUrl/$($documentUrls.ApiUrl)"
377+
$url = "$baseUrl/$($documentUrls.ApiUrl)"
377378

378-
$now = Get-Time
379+
$now = Get-Time
379380

380-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $documentUrls.ResourceUrl -now $now
381+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $documentUrls.ResourceUrl -now $now
381382

382-
$requestPartitionKey = if ($PartitionKey) { $PartitionKey } else { $RecordId }
383-
}
384-
process {
385-
try {
386-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey -isQuery $true
383+
$requestPartitionKey = if ($PartitionKey) { $PartitionKey } else { $RecordId }
387384

388-
Invoke-CosmosDbApiRequest -Verb $GET_VERB -Url $url -Headers $headers
389-
}
390-
catch {
391-
Get-ExceptionResponseOrThrow $_
392-
}
385+
try {
386+
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -PartitionKey $requestPartitionKey -isQuery $true
387+
388+
Invoke-CosmosDbApiRequest -Verb $GET_VERB -Url $url -Headers $headers
393389
}
390+
catch {
391+
Get-ExceptionResponseOrThrow $_
392+
}
394393
}
395394

396395
<#
@@ -443,31 +442,29 @@ Function Get-AllCosmosDbRecords(
443442
[parameter(Mandatory = $true)][string]$Database,
444443
[parameter(Mandatory = $true)][string]$Container,
445444
[parameter(Mandatory = $true)][string]$Collection,
446-
[parameter(Mandatory = $false)][string]$SubscriptionId = "") {
447-
begin {
448-
$baseUrl = Get-BaseDatabaseUrl $Database
449-
$collectionsUrl = Get-CollectionsUrl $Container $Collection
450-
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
445+
[parameter(Mandatory = $false)][string]$SubscriptionId = ""
446+
) {
447+
$baseUrl = Get-BaseDatabaseUrl $Database
448+
$collectionsUrl = Get-CollectionsUrl $Container $Collection
449+
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
451450

452-
$url = "$baseUrl/$docsUrl"
451+
$url = "$baseUrl/$docsUrl"
453452

454-
$now = Get-Time
453+
$now = Get-Time
455454

456-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
457-
}
458-
process {
459-
$tmp = $ProgressPreference
460-
$ProgressPreference = 'SilentlyContinue'
461-
try {
462-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true
455+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $GET_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
463456

464-
Invoke-CosmosDbApiRequestWithContinuation -verb $GET_VERB -url $url -Headers $headers
465-
}
466-
catch {
467-
Get-ExceptionResponseOrThrow $_
468-
}
469-
$ProgressPreference = $tmp
457+
$tmp = $ProgressPreference
458+
$ProgressPreference = 'SilentlyContinue'
459+
try {
460+
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true
461+
462+
Invoke-CosmosDbApiRequestWithContinuation -verb $GET_VERB -url $url -Headers $headers
470463
}
464+
catch {
465+
Get-ExceptionResponseOrThrow $_
466+
}
467+
$ProgressPreference = $tmp
471468
}
472469

473470
<#
@@ -541,39 +538,37 @@ Function Search-CosmosDbRecords(
541538
[parameter(Mandatory = $true)][string]$Query,
542539
[parameter(Mandatory = $false)]$Parameters = $null,
543540
[parameter(Mandatory = $false)][string]$SubscriptionId = "",
544-
[parameter(Mandatory = $false)][switch]$DisableExtraFeatures = $false) {
545-
begin {
546-
$Parameters = @(Get-QueryParametersAsNameValuePairs $Parameters)
541+
[parameter(Mandatory = $false)][switch]$DisableExtraFeatures = $false
542+
) {
543+
$Parameters = @(Get-QueryParametersAsNameValuePairs $Parameters)
547544

548-
$baseUrl = Get-BaseDatabaseUrl $Database
549-
$collectionsUrl = Get-CollectionsUrl $Container $Collection
550-
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
545+
$baseUrl = Get-BaseDatabaseUrl $Database
546+
$collectionsUrl = Get-CollectionsUrl $Container $Collection
547+
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
551548

552-
$url = "$baseUrl/$docsUrl"
549+
$url = "$baseUrl/$docsUrl"
553550

554-
$now = Get-Time
551+
$now = Get-Time
555552

556-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
553+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
554+
555+
if (!$DisableExtraFeatures) {
556+
return Search-CosmosDbRecordsWithExtraFeatures -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -Query $Query -Parameters $Parameters -SubscriptionId $SubscriptionId
557557
}
558-
process {
559-
if (!$DisableExtraFeatures) {
560-
return Search-CosmosDbRecordsWithExtraFeatures -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -Query $Query -Parameters $Parameters -SubscriptionId $SubscriptionId
561-
}
562558

563-
try {
564-
$body = @{
565-
query = $Query;
566-
parameters = $Parameters;
567-
}
559+
try {
560+
$body = @{
561+
query = $Query;
562+
parameters = $Parameters;
563+
}
568564

569-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
570-
$headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
565+
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
566+
$headers["x-ms-documentdb-query-enablecrosspartition"] = "true"
571567

572-
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
573-
}
574-
catch {
575-
Get-ExceptionResponseOrThrow $_
576-
}
568+
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
569+
}
570+
catch {
571+
Get-ExceptionResponseOrThrow $_
577572
}
578573
}
579574

@@ -587,73 +582,70 @@ Function Search-CosmosDbRecordsWithExtraFeatures
587582
$Parameters,
588583
[string]$SubscriptionId
589584
) {
590-
begin {
591-
$Parameters = @(Get-QueryParametersAsNameValuePairs $Parameters)
585+
$Parameters = @(Get-QueryParametersAsNameValuePairs $Parameters)
592586

593-
$baseUrl = Get-BaseDatabaseUrl $Database
594-
$collectionsUrl = Get-CollectionsUrl $Container $Collection
595-
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
587+
$baseUrl = Get-BaseDatabaseUrl $Database
588+
$collectionsUrl = Get-CollectionsUrl $Container $Collection
589+
$docsUrl = "$collectionsUrl/$DOCS_TYPE"
596590

597-
$url = "$baseUrl/$docsUrl"
591+
$url = "$baseUrl/$docsUrl"
598592

599-
$now = Get-Time
593+
$now = Get-Time
600594

601-
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
595+
$encodedAuthString = Get-AuthorizationHeader -ResourceGroup $ResourceGroup -SubscriptionId $SubscriptionId -Database $Database -verb $POST_VERB -resourceType $DOCS_TYPE -resourceUrl $collectionsUrl -now $now
596+
597+
$allPartitionKeyRangesOrError = Get-PartitionKeyRangesOrError -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -SubscriptionId $SubscriptionId
602598

603-
$allPartitionKeyRangesOrError = Get-PartitionKeyRangesOrError -ResourceGroup $ResourceGroup -Database $Database -Container $Container -Collection $Collection -SubscriptionId $SubscriptionId
599+
if ($allPartitionKeyRangesOrError.ErrorRecord) {
600+
return Get-ExceptionResponseOrThrow $allPartitionKeyRangesOrError.ErrorRecord
604601
}
605-
process {
606-
if ($allPartitionKeyRangesOrError.ErrorRecord) {
607-
return Get-ExceptionResponseOrThrow $allPartitionKeyRangesOrError.ErrorRecord
608-
}
609-
610-
try {
611-
$ranges = $allPartitionKeyRangesOrError.Ranges
602+
603+
try {
604+
$ranges = $allPartitionKeyRangesOrError.Ranges
612605

613-
$body = @{
614-
query = $Query;
615-
parameters = $Parameters;
616-
}
606+
$body = @{
607+
query = $Query;
608+
parameters = $Parameters;
609+
}
617610

618-
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
619-
$headers += @{
620-
"x-ms-documentdb-query-enablecrosspartition" = "true";
621-
"x-ms-cosmos-supported-query-features" = "NonValueAggregate, Aggregate, Distinct, MultipleOrderBy, OffsetAndLimit, OrderBy, Top, CompositeAggregate, GroupBy, MultipleAggregates";
622-
"x-ms-documentdb-query-enable-scan" = "true";
623-
"x-ms-documentdb-query-parallelizecrosspartitionquery" = "true";
624-
"x-ms-cosmos-is-query-plan-request" = "True";
625-
}
611+
$headers = Get-CommonHeaders -now $now -encodedAuthString $encodedAuthString -isQuery $true -contentType "application/Query+json"
612+
$headers += @{
613+
"x-ms-documentdb-query-enablecrosspartition" = "true";
614+
"x-ms-cosmos-supported-query-features" = "NonValueAggregate, Aggregate, Distinct, MultipleOrderBy, OffsetAndLimit, OrderBy, Top, CompositeAggregate, GroupBy, MultipleAggregates";
615+
"x-ms-documentdb-query-enable-scan" = "true";
616+
"x-ms-documentdb-query-parallelizecrosspartitionquery" = "true";
617+
"x-ms-cosmos-is-query-plan-request" = "True";
618+
}
626619

627-
$queryPlan = Invoke-CosmosDbApiRequest -verb $POST_VERB -url $url -Body $body -Headers $headers | Get-CosmosDbRecordContent
628-
629-
$rewrittenQuery = $queryPlan.QueryInfo.RewrittenQuery
630-
$searchQuery = if ($rewrittenQuery) { $rewrittenQuery } else { $Query };
620+
$queryPlan = Invoke-CosmosDbApiRequest -verb $POST_VERB -url $url -Body $body -Headers $headers | Get-CosmosDbRecordContent
621+
622+
$rewrittenQuery = $queryPlan.QueryInfo.RewrittenQuery
623+
$searchQuery = if ($rewrittenQuery) { $rewrittenQuery } else { $Query };
631624

632-
$body = @{
633-
query = $searchQuery;
634-
parameters = $Parameters;
635-
}
625+
$body = @{
626+
query = $searchQuery;
627+
parameters = $Parameters;
628+
}
636629

637-
$headers.Remove("x-ms-cosmos-is-query-plan-request")
630+
$headers.Remove("x-ms-cosmos-is-query-plan-request")
638631

639-
$partitionKeyRanges =
640-
if ($env:COSMOS_DB_FLAG_ENABLE_PARTITION_KEY_RANGE_SEARCHES -eq 1) {
641-
Get-FilteredPartitionKeyRangesForQuery -AllRanges $ranges -QueryRanges $queryPlan.QueryRanges
642-
}
643-
else {
644-
$ranges
645-
}
632+
$partitionKeyRanges =
633+
if ($env:COSMOS_DB_FLAG_ENABLE_PARTITION_KEY_RANGE_SEARCHES -eq 1) {
634+
Get-FilteredPartitionKeyRangesForQuery -AllRanges $ranges -QueryRanges $queryPlan.QueryRanges
635+
}
636+
else {
637+
$ranges
638+
}
646639

647-
foreach ($partitionKeyRange in $partitionKeyRanges) {
648-
$headers["x-ms-documentdb-partitionkeyrangeid"] = $partitionKeyRange.id
640+
foreach ($partitionKeyRange in $partitionKeyRanges) {
641+
$headers["x-ms-documentdb-partitionkeyrangeid"] = $partitionKeyRange.id
649642

650-
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
651-
}
643+
Invoke-CosmosDbApiRequestWithContinuation -verb $POST_VERB -url $url -Body $body -Headers $headers
652644
}
653-
catch {
654-
Get-ExceptionResponseOrThrow $_
655-
}
656645
}
646+
catch {
647+
Get-ExceptionResponseOrThrow $_
648+
}
657649
}
658650

659651
<#
@@ -800,7 +792,6 @@ Function Update-CosmosDbRecord {
800792
[parameter(Mandatory = $false, ParameterSetName = "ExplicitPartitionKey")][string]$PartitionKey = "",
801793
[parameter(Mandatory = $false, ParameterSetName = "ParttionKeyCallback")]$GetPartitionKeyBlock = $null,
802794
[parameter(Mandatory = $false)][bool]$EnforceOptimisticConcurrency = $true
803-
804795
)
805796

806797
begin {
@@ -947,21 +938,23 @@ Function Get-CosmosDbRecordContent([parameter(ValueFromPipeline)]$RecordResponse
947938
}
948939
elseif ($code -eq 401) {
949940
if ($env:COSMOS_DB_FLAG_ENABLE_READONLY_KEYS -eq 1) {
950-
throw "Unauthorized (used a readonly key)"
941+
throw "(401) Unauthorized (used a readonly key)"
951942
}
952-
throw "Unauthorized"
943+
throw "(401) Unauthorized"
953944
}
954945
elseif ($code -eq 404) {
955946
if ($content.Message -like "*Owner resource does not exist*") {
956-
throw "Database does not exist"
947+
throw "(404) Database does not exist"
957948
}
958949

959-
throw "Record not found"
950+
throw "(404) Record not found"
951+
}
952+
elseif ($code -eq 412) {
953+
throw "(412) Conflict"
960954
}
961955
elseif ($code -eq 429) {
962-
throw "Request rate limited"
956+
throw "(429) Request rate limited"
963957
}
964-
965958
else {
966959
$message = $content.Message
967960
throw "Request failed with status code $code with message`n`n$message"

0 commit comments

Comments
 (0)