Fetch Exchange mailboxes via REST API in parallel using Async and HttpClient

ExoHttpClient.ps1

using namespace Microsoft.Exchange.Management.AdminApiProvider
using namespace Microsoft.Exchange.Management.ExoPowershellSnapin
using namespace System.Collections.Generic
using namespace System.Net.Http
using namespace System.Net.Http.Headers
using namespace System.Threading.

<#
I made this as a way to overcome some limitations with very complicated filter queries that exceed the REST limit, while being able to run them in parallel without resorting to runspaces
#>


#This static type will be "super global" and easily accessible across runspaces and allow us to build multiple httpclients
#Technically you're supposed to only use one but I like the convenience of DefaultRequestHeaders for the auth token
try {
  #Test if the type exists
  [void][ExoHttpClient]
} catch {
  Add-Type -TypeDefinition @'
    using System.Collections.Generic;
    using System.Net.Http;
    public class ExoHttpClient {
      public static Dictionary<string,HttpClient> Client = new Dictionary<string,HttpClient>();
    }
'@
}

function Get-ExoHttpClient {
  <#
  .SYNOPSIS fetches an HTTPClient for the current exchange connection. Multiple connections to different tenants can be made and fetched separately
  #>
  [CmdletBinding()]
  param (
    [guid]$ContextId = (Get-ExoCurrentConnectionId),
    [int]$BatchSize = 1000
  )

  $client = [ExoHttpClient]::Client[([string]$ContextId)]

  if (-not $client) {
    $newClient = [HttpClient]::new()
    #This appears to be a hardcoded guid for the API
    [guid]$apiGuid = 'a9a9ae2a-2e0f-47ac-bdaa-0e8fa458a60e'
    $newClient.BaseAddress = "https://outlook.office365.com/adminapi/beta/$apiGuid/"
    $newClient.Timeout = '00:00:30'
    $newClient.DefaultRequestHeaders.Add('Prefer', "odata.maxpagesize=$BatchSize")
    [ExoHttpClient]::Client[$ContextId] = $newClient
    $client = $newClient
  }

  #Refresh the token if required
  $context = [ConnectionContextFactory]::GetCurrentConnectionContext($ContextId)
  $tokenProvider = $context.TokenProvider
  $token = $tokenProvider.GetValidTokenFromCache()
  $scheme, $token = $token.AuthorizationHeader -split ' '
  $client.DefaultRequestHeaders.Authorization = [AuthenticationHeaderValue]::new(
    <#Scheme#>    $scheme,
    <#Parameter#> $token
  )

  return $client
}

function Get-ExoCurrentConnectionId {
  #First get the module path
  [string]$exchRestModulePath = Get-Module tmp*
  | Where-Object Description -EQ 'This is a Powershell module generated by using the AutoGEN infra.' #HACK: Need better way to identify
  | Select-Object -Last 1 #Multiple modules can be created, get the latest one
  | ForEach-Object Path

  if (-not $exchRestModulePath) {
    throw 'You must connect using Connect-ExchangeOnline first to use this command.'
  }

  #Then extract the connection context
  if ((Get-Content -Raw $exchRestModulePath) -notmatch "ConnectionContextObjectId = '([\w-]+)'") {
    throw 'Exchange module found but connection context missing. This is a bug or you are not on 2.0.6+'
  }
  [guid]$ContextId = $matches[1]
  return $contextId
}

enum ExoRecipientType {
  UserMailbox
  MailUser
  MailContact
  DynamicDistributionGroup
  MailUniversalDistributionGroup
  MailUniversalSecurityGroup
  MailNonUniversalGroup
  PublicFolder
}

function Get-ExoRestRecipient {
  [CmdletBinding()]
  param(
    [Parameter(Mandatory, ValueFromPipeline)][string]$Filter,
    [ExoRecipientType]$RecipientType,
    [HttpClient]$client = $(Get-ExoHttpClient)
  )

  process {
    $query = ConvertTo-QueryString @{
      '$filter' = $Filter
    }

    $task = $client.GetStringAsync("Recipient?$query")

    #This will be used by the receive-exorestrecipient
    [KeyValuePair[Task[String], string]]::new($task, $filter)
  }
}

function Receive-ExoRestRecipient {
  [CmdletBinding()]
  param(
    [Parameter(ValueFromPipeline)][KeyValuePair[Task[String], string]]$RecipientRequest,
    #TODO: Bring httpclient along for the ride
    [HttpClient]$client = $(Get-ExoHttpClient)
  )
  begin {
    #We will use this to map tasks to their filters for reference
    [Dictionary[Task[String], string]]$tasks = @{}
  }
  process {
    #Fill up the queue
    $tasks.Add($RecipientRequest.key, $RecipientRequest.value)
  }
  end {
    while ($tasks.Count -ne 0) {
      $taskList = @($tasks.Keys)
      $doneTaskIndex = [Task]::WaitAny($taskList)
      $doneTask = $taskList[$doneTaskIndex]
      $taskName = $tasks[$doneTask]
      [void]$tasks.Remove($doneTask)
      Write-Host -fore Cyan ('Task {0} {1} - {2}' -f $doneTask.Id, $doneTask.Status, $taskName)
      try {
        $result = $doneTask.GetAwaiter().GetResult()
        | ConvertFrom-Json
      } catch {
        Write-Warning "$PSItem"
        return
      }

      $skipLink = $result.'@odata.nextLink'
      if ($skipLink) {
        Write-Host -fore DarkCyan "Adding SkipLink task: $skipLink"
        #Get the next batch
        $tasks.Add($client.GetStringAsync($skipLink), $skipLink)
      }

      [Recipient[]]$returnData = $result
      | ForEach-Object Value
      | Select-Object -exclude '@odata.id', '@odata.editlink'

      $returnData
    }
  }
}

function ConvertTo-QueryString {
  [CmdletBinding()]
  param(
    [Parameter(Mandatory, ValueFromPipeline)][hashtable]$hashtable
  )
  #This is the only way I know to instantiate this simple collection that can parse to a query string
  $queryCollection = [System.Web.HttpUtility]::ParseQueryString('')

  $hashtable.GetEnumerator() | ForEach-Object {
    $queryCollection.Add($_.Key, $_.Value)
  }
  $queryCollection.ToString()
}

# Example
#Connect-ExchangeOnline
#$filters[0..20] | Get-ExoRestRecipient | Receive-ExoRestRecipient