IntuneDiff

1.1.1

Private/Graph/Invoke-GraphBatchRequest.ps1

# Copyright (c) 2026 Sandy Zeng. All rights reserved.
# Source-available. All rights reserved. See LICENSE file.

<#
    Invoke-GraphBatchRequest.ps1 — Sends batched Microsoft Graph API requests
    (up to 20 per call) via Invoke-RestMethod with 429 + 401 handling.
 
    Author: Sandy Zeng
    Project: IntuneDiff
 
    Version History:
    1.0.0 Initial release.
    2.0.0 Rewritten to use Invoke-RestMethod with a Bearer token — no Microsoft.Graph SDK.
#>

function Invoke-GraphBatchRequest {
    <#
    .SYNOPSIS
        Sends a batch of up to 20 Graph API requests in a single POST /$batch call.
 
    .OUTPUTS
        Hashtable keyed by request id. Each value is the response body or $null on failure.
    #>
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [object[]]$Requests,

        [int]$MaxRetries = 3
    )

    $batchSize = 20
    $results   = @{}

    $chunks = [System.Collections.Generic.List[object[]]]::new()
    for ($i = 0; $i -lt $Requests.Count; $i += $batchSize) {
        $end = [Math]::Min($i + $batchSize, $Requests.Count)
        $chunks.Add($Requests[$i..($end - 1)])
    }

    foreach ($chunk in $chunks) {
        $pendingRequests = @($chunk)
        $retryCount      = 0
        $tokenRefreshed  = $false

        while ($pendingRequests.Count -gt 0 -and $retryCount -le $MaxRetries) {
            $batchBody = @{
                requests = @($pendingRequests | ForEach-Object {
                    @{ id = [string]$_.id; method = $_.method; url = $_.url }
                })
            }

            Write-IDLog "Graph BATCH: $($pendingRequests.Count) requests (attempt $($retryCount + 1))"

            try {
                $jsonBody = $batchBody | ConvertTo-Json -Depth 10 -Compress
                $headers  = @{
                    Authorization = "Bearer $(Get-IDAccessToken)"
                    Accept        = 'application/json'
                }
                $response = Invoke-RestMethod -Method POST `
                    -Uri 'https://graph.microsoft.com/beta/$batch' `
                    -Headers $headers `
                    -Body $jsonBody `
                    -ContentType 'application/json' `
                    -TimeoutSec 60 `
                    -ErrorAction Stop
            } catch {
                $statusCode = $null
                if ($_.Exception.Response) { $statusCode = [int]$_.Exception.Response.StatusCode }
                Write-IDLog "Batch POST failed (status $statusCode): $($_.Exception.Message)"

                if ($statusCode -eq 401 -and -not $tokenRefreshed) {
                    Write-IDLog 'Batch 401 — refreshing token and retrying'
                    try {
                        Get-IDAccessToken -ForceRefresh | Out-Null
                        $tokenRefreshed = $true
                        $retryCount++
                        continue
                    } catch {
                        Write-IDLog "Token refresh failed: $($_.Exception.Message)"
                    }
                }

                foreach ($req in $pendingRequests) { $results[[string]$req.id] = $null }
                break
            }

            $throttled     = [System.Collections.Generic.List[object]]::new()
            $maxRetryAfter = 0

            # Invoke-RestMethod with JSON yields PSCustomObject; old code also accepted IDictionary
            $responses = $null
            if ($response -is [System.Collections.IDictionary]) {
                $responses = $response['responses']
            } elseif ($response.PSObject.Properties['responses']) {
                $responses = $response.responses
            }

            if (-not $responses -or @($responses).Count -eq 0) {
                Write-IDLog "Batch response had no 'responses' array. Marking all as failed."
                foreach ($req in $pendingRequests) { $results[[string]$req.id] = $null }
                break
            }

            foreach ($resp in $responses) {
                if ($resp -is [System.Collections.IDictionary]) {
                    $id     = [string]$resp['id']
                    $status = [int]$resp['status']
                    $body   = $resp['body']
                    $hdrs   = $resp['headers']
                } else {
                    $id     = [string]$resp.id
                    $status = [int]$resp.status
                    $body   = $resp.body
                    $hdrs   = $resp.headers
                }

                if ($status -ge 200 -and $status -lt 300) {
                    $results[$id] = $body
                } elseif ($status -eq 429) {
                    $retryAfter = 2
                    if ($hdrs) {
                        $ra = if ($hdrs -is [System.Collections.IDictionary]) { $hdrs['Retry-After'] } else { $hdrs.'Retry-After' }
                        if ($ra) { try { $retryAfter = [int]$ra } catch {} }
                    }
                    if ($retryAfter -gt $maxRetryAfter) { $maxRetryAfter = $retryAfter }
                    $original = $pendingRequests | Where-Object { [string]$_.id -eq $id } | Select-Object -First 1
                    if ($original) { $throttled.Add($original) }
                } elseif ($status -eq 401 -and -not $tokenRefreshed) {
                    # Per-item 401: refresh token, retry whole batch once
                    Write-IDLog "Batch item $id returned 401 — refreshing token"
                    try {
                        Get-IDAccessToken -ForceRefresh | Out-Null
                        $tokenRefreshed = $true
                    } catch {
                        Write-IDLog "Token refresh failed: $($_.Exception.Message)"
                    }
                    $original = $pendingRequests | Where-Object { [string]$_.id -eq $id } | Select-Object -First 1
                    if ($original) { $throttled.Add($original) }
                } else {
                    Write-IDLog "Batch item $id failed with status $status"
                    $results[$id] = $null
                }
            }

            if ($throttled.Count -gt 0 -and $retryCount -lt $MaxRetries) {
                if ($maxRetryAfter -gt 0) {
                    Write-IDLog "Batch: $($throttled.Count) items retrying after ${maxRetryAfter}s..."
                    Start-Sleep -Seconds ([Math]::Max($maxRetryAfter, 1))
                }
                $pendingRequests = @($throttled.ToArray())
                $retryCount++
            } else {
                foreach ($req in $throttled) { $results[[string]$req.id] = $null }
                break
            }
        }
    }

    return $results
}