Merge pull request #7 from cmbrose/continuation-token-fix

Fix continuation tokens on PS 7+

Author: cmbrose

CHANGELOG.md

@@ -1,6 +1,9 @@
 # Release Notes
 All notable changes and release history of the "cosmos-db" module will be documented in this file.
 
+## 1.8
+* Bugfix for continuation tokens on Powershell version 7+
+
 ## 1.7
 * Minor interface improvements
 

cosmos-db/cosmos-db.psd1

@@ -3,16 +3,15 @@
 #
 # Generated by: cmbrose
 #
-# Generated on: 12/27/2020
-#
+# Generated on: 1/1/2021
 
 @{
 
 # Script module or binary module file associated with this manifest.
 # RootModule = ''
 
 # Version number of this module.
-ModuleVersion = '1.7'
+ModuleVersion = '1.8'
 
 # Supported PSEditions
 # CompatiblePSEditions = @()

cosmos-db/cosmos-db.psm1

@@ -187,19 +187,43 @@ Function Invoke-CosmosDbApiRequest([string]$verb, [string]$url, $headers, $body=
     Invoke-WebRequest -Method $verb -Uri $url -Body $body -Headers $headers
 }
 
+Function Get-ContinuationToken($response)
+{
+    $value = $response.Headers["x-ms-continuation"]
+
+    if ($PSVersionTable.PSEdition -eq "Core")
+    {
+        if (-not $value)
+        {
+            return $null
+        }
+
+        # Headers were changed to arrays in version 7
+        # https://docs.microsoft.com/en-us/powershell/scripting/whats-new/breaking-changes-ps6?view=powershell-7.1#changes-to-web-cmdlets
+        $value[0]
+    }
+    else
+    {
+        $value
+    }
+}
+
 Function Invoke-CosmosDbApiRequestWithContinuation([string]$verb, [string]$url, $headers, $body=$null)
 {
     process
     {
         $response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
         $response
 
-        while ($response.Headers["x-ms-continuation"])
+        $continuationToken = Get-ContinuationToken $response
+        while ($continuationToken)
         {
-            $headers["x-ms-continuation"] = $response.Headers["x-ms-continuation"];
+            $headers["x-ms-continuation"] = $continuationToken
 
             $response = Invoke-CosmosDbApiRequest -Verb $verb -Url $url -Body $body -Headers $headers
             $response
+
+            $continuationToken = Get-ContinuationToken $response
         }   
     }
 }

tests/Invoke-CosmosDbApiRequestWithContinuation.Tests.ps1

@@ -6,6 +6,12 @@ InModuleScope cosmos-db {
         Use-CosmosDbInternalFlag -EnableCaching $false
 
         . $PSScriptRoot\Utils.ps1
+
+        $ORIG_PS_EDITION = $PSVersionTable.PSEdition
+    }
+
+    AfterAll {
+        $PSVersionTable.PSEdition = $ORIG_PS_EDITION
     }
 
     Describe "Invoke-CosmosDbApiRequestWithContinuation" {
@@ -46,6 +52,45 @@ InModuleScope cosmos-db {
             Assert-MockCalled Invoke-CosmosDbApiRequest -Times 1
         }
 
+        It "Handles responses without continuation header for Powershell 7" {  
+            $PSVersionTable.PSEdition = "Core"
+
+            $MOCK_VERB = "MOCK_VERB"
+            $MOCK_URL = "MOCK_URL"
+            $MOCK_BODY = @{
+                Mock = "Mock"
+            }
+            $MOCK_HEADERS = @{
+                Mock = "Mock";
+            }
+
+            $response = @{
+                StatusCode = 200;
+                Content = "{}";
+                Headers = @{};
+            }
+
+            Mock Invoke-CosmosDbApiRequest {
+                param($verb, $url, $body, $headers) 
+
+                $headers["x-ms-continuation"] | Should -BeNullOrEmpty |  Out-Null
+
+                $verb | Should -Be $MOCK_VERB | Out-Null
+                $url | Should -Be $MOCK_URL | Out-Null
+                $body | Should -Be $MOCK_BODY | Out-Null
+                $headers | Should -Be $MOCK_HEADERS | Out-Null    
+
+                $response
+            }
+
+            $result = Invoke-CosmosDbApiRequestWithContinuation -Verb $MOCK_VERB -Url $MOCK_URL -Body $MOCK_BODY -Headers $MOCK_HEADERS
+
+            $result | Should -BeExactly $response
+            @($result).Count | Should -Be 1
+
+            Assert-MockCalled Invoke-CosmosDbApiRequest -Times 1
+        }
+
         It "Handles continuation response headers" {  
             $continuationTokens = @($null, "1", "2", "3")
 
@@ -92,5 +137,56 @@ InModuleScope cosmos-db {
 
             Assert-MockCalled Invoke-CosmosDbApiRequest -Times $continuationTokens.Count
         }
+
+        It "Handles continuation response headers for Powershell 7" {  
+            $continuationTokens = @($null, "100", "200", "300")
+
+            $PSVersionTable.PSEdition = "Core"
+
+            $MOCK_VERB = "MOCK_VERB"
+            $MOCK_URL = "MOCK_URL"
+            $MOCK_BODY = @{
+                Mock = "Mock"
+            }
+            $MOCK_HEADERS = @{
+                Mock = "Mock";
+            }
+
+            $global:idx = 0
+            $global:expectedResponses = @()
+
+            Mock Invoke-CosmosDbApiRequest {
+                param($verb, $url, $body, $headers)             
+
+                $verb | Should -Be $MOCK_VERB | Out-Null
+                $url | Should -Be $MOCK_URL | Out-Null
+                $body | Should -Be $MOCK_BODY | Out-Null
+
+                $headers["x-ms-continuation"] | Should -Be $continuationTokens[$global:idx] | Out-Null
+                $global:idx = $global:idx + 1
+
+                $headers.Remove("x-ms-continuation")
+                AssertHashtablesEqual $MOCK_HEADERS $headers
+        
+                $response = @{
+                    StatusCode = 200;
+                    Content = "$global:idx";
+                    Headers = @{
+                        # The empty here is to trick powershell into no automatically converting the single item array into just a value
+                        "x-ms-continuation" = @($continuationTokens[$global:idx], "")
+                    };
+                }
+                
+                $global:expectedResponses += $response
+                $response
+            }
+
+            $result = Invoke-CosmosDbApiRequestWithContinuation -Verb $MOCK_VERB -Url $MOCK_URL -Body $MOCK_BODY -Headers $MOCK_HEADERS
+
+            $result | Should -BeExactly $global:expectedResponses
+            @($result).Count | Should -Be $continuationTokens.Count
+
+            Assert-MockCalled Invoke-CosmosDbApiRequest -Times $continuationTokens.Count
+        }
     }
 }