Rate-Limiting, Powershell, Pester, and the ZenDesk API

Have you ever had to utterly hammer an API via a little Powershell script, only to find you’re getting rate limited, and lose all that previously downloaded data before you could persist it somewhere?

I have. I’ve recently put together a script to query a customer support ticketing system’s API, get a list of all tickets within a given time period, then query again to get to get the full content for each ticket.

All of this data is being added to a custom Powershell object (as opposed to churning out to a file as I go) since I convert it all to Json right at the end.

I’d rather not get half way through a few thousand calls and then have the process fail due to being throttled, so I’ve chosen to create a little function that will check for – in my case – an HTTP 429 response (“Too Many Requests”), get the value of the “Retry-After” header, then wait that many seconds before trying again.

This particular implementation is all quite specific to the ZenDesk API, but could easily be adapted to other APIs which rate limit/throttle and return the appropriate headers.

I’ve called the function CalmlyCall as a nod to Twitter’s original throttling implementation when they used HTTP 420 – “Enhance Your Calm” – back in the day.

function CalmlyCall($uri, $headers) {
    $enhanceyourcalm = 0

    do {
        if ($enhanceyourcalm -gt 0) {
            Write-Host -Object "Throttled; calming down for $enhanceyourcalm seconds..."
            Start-Sleep -Seconds $enhanceyourcalm
            $headers.Calm = $true
        }

        $response = Invoke-WebRequest -Uri $uri -Headers $headers -UseBasicParsing
        $ratelimit = $response.Headers['X-Rate-Limit-Remaining']

        if ([int]$ratelimit -lt 50) {
            Write-Host -Object "Calm may be needed: rate limit remaining - $($ratelimit)"
        }

        $enhanceyourcalm = $response.Headers["Retry-After"]
    } while ($response.StatusCode -eq 429)

    return ConvertFrom-Json $response.Content
}

You’ll see that I’m just using a do .. while loop to call the specified endpoint, check the headers for throttling information, and patiently (using Start-Sleep) until I’m allowed to continue. Simple stuff.

I’m only adding a Calm header when I’m throttled to make it more easily testable. I’m sure there’s a better solution to this, but it works well enough for me.

And since I’m a fan of testing my powershell, here’s the associated Pester script – it’s not great, since I’m mocking against Write-Host, for example:

$here = Split-Path -Parent $MyInvocation.MyCommand.Path
$sut = (Split-Path -Leaf $MyInvocation.MyCommand.Path) -replace '\.Tests\.', '.'
. "$here\$sut"

Describe "CalmlyCall" {
    Context "if no rate limiting exists" {
        It "should call the endpoint" {
            Mock Write-Host {}
            Mock Invoke-WebRequest -ParameterFilter { $Uri.ToString().EndsWith("test.here") } {        
                [PSCustomObject] `
                @{
                    Headers = @{};
                    StatusCode = 200;
                    Content = "{'result':'mycontent'}"
                }
            }

            $actual = CalmlyCall -uri "test.here" -headers @{}

            Assert-MockCalled Invoke-WebRequest
            $actual.result | Should be "mycontent"
        }
    }


    Context "if rate limiting exists" {
        It "should check limit remaining" {
            Mock Invoke-WebRequest -ParameterFilter { $Uri.ToString().EndsWith("test.here") } {        
                [PSCustomObject] `
                @{
                    Headers = @{'X-Rate-Limit-Remaining' = '20';'Retry-After' = '0'};
                    StatusCode = 200;
                    Content = "{'result':'mycontent'}"
                }
            }

            Mock Write-Host -ParameterFilter { $Object.ToString().EndsWith("rate limit remaining - 20")}

            $actual = CalmlyCall -uri "test.here" -headers @{}

            Assert-MockCalled Invoke-WebRequest
            Assert-MockCalled Write-Host -Times 1 -Exactly
            $actual.result | Should be "mycontent"
        }
    }


    Context "if initial request is throttled" {
        It "should respect throttling" {
            Mock Invoke-WebRequest -ParameterFilter { $Uri.ToString().EndsWith("test.here") } {        
                [PSCustomObject] `
                @{
                    Headers = @{'X-Rate-Limit-Remaining' = '0';'Retry-After' = '10'};
                    StatusCode = 429;
                    Content = "{'result':'mycontent'}"
                }
            }

            Mock Invoke-WebRequest -ParameterFilter { $Uri.ToString().EndsWith("test.here") -and $Headers -ne $null -and $Headers.Calm} {        
                [PSCustomObject] `
                @{
                    Headers = @{'X-Rate-Limit-Remaining' = '100';'Retry-After' = '0'};
                    StatusCode = 200;
                    Content = "{'result':'mycontent'}"
                }
            }

            Mock Start-Sleep -ParameterFilter { $Seconds -eq 10 }
            Mock Write-Host -ParameterFilter { $Object.ToString().EndsWith("rate limit remaining - 0") -or
                                               $Object.ToString().StartsWith("Throttled; calming down for 10 seconds")}

            $actual = CalmlyCall -uri "test.here" -headers @{}

            Assert-MockCalled Invoke-WebRequest -Times 2 -Exactly
            Assert-MockCalled Start-Sleep -Times 1 -Exactly
            Assert-MockCalled Write-Host -Times 2 -Exactly
            $actual.result | Should be "mycontent"
        }
    }
}

Feel free to suggest improvements!

Leave a Reply

Your email address will not be published. Required fields are marked *