From 8d65d84c898f13df7ce9008b265eb3c8ee9dfc20 Mon Sep 17 00:00:00 2001
From: gregslack78 <85254974+gregslack78@users.noreply.github.com>
Date: Sun, 2 Feb 2025 18:57:56 -0500
Subject: [PATCH] Updates to use dynamic pagination (#868)

---
 .../Helpers/Get-AzPolicyOrSetDefinitions.ps1  |  4 +-
 Scripts/Helpers/Search-AzGraphAllItems.ps1    | 45 +++++++++++++++----
 2 files changed, 40 insertions(+), 9 deletions(-)

diff --git a/Scripts/Helpers/Get-AzPolicyOrSetDefinitions.ps1 b/Scripts/Helpers/Get-AzPolicyOrSetDefinitions.ps1
index 5009d22f..4d297dfb 100644
--- a/Scripts/Helpers/Get-AzPolicyOrSetDefinitions.ps1
+++ b/Scripts/Helpers/Get-AzPolicyOrSetDefinitions.ps1
@@ -26,15 +26,17 @@ function Get-AzPolicyOrSetDefinitions {
             $query = "PolicyResources | where type == 'microsoft.authorization/policydefinitions'"
             $progressItemName = "Policy definitions"
             $excludedIds = $desiredState.excludedPolicyDefinitions
+            $progressIncrement = 1000
         }
         policySetDefinitions {
             $query = "PolicyResources | where type == 'microsoft.authorization/policysetdefinitions'"
             $progressItemName = "Policy Set definitions"
             $excludedIds = $desiredState.excludedPolicySetDefinitions
+            $progressIncrement = 250
         }
     }
 
-    $policyResources = Search-AzGraphAllItems -Query $query -ProgressItemName $progressItemName
+    $policyResources = Search-AzGraphAllItems -Query $query -ProgressItemName $progressItemName -ProgressIncrement $progressIncrement
     foreach ($policyResource in $policyResources) {
         $resourceTenantId = $policyResource.tenantId
         if ($resourceTenantId -in @($null, "", $environmentTenantId)) {
diff --git a/Scripts/Helpers/Search-AzGraphAllItems.ps1 b/Scripts/Helpers/Search-AzGraphAllItems.ps1
index 2a13cccf..0b800d2b 100644
--- a/Scripts/Helpers/Search-AzGraphAllItems.ps1
+++ b/Scripts/Helpers/Search-AzGraphAllItems.ps1
@@ -7,12 +7,19 @@ function Search-AzGraphAllItems {
     )
 
     # Search-AzGraph can only return a maximum of 1000 items. Without the -First it will only return 100 items
-    $body = @{
-        query = $Query
-        # options = @{
-        #     "`$top"  = 1000
-        #     "`$skip" = 0
-        # }
+    if ($ProgressItemName -ne "Policy Set definitions") {
+        $body = @{
+            query = $Query
+        }
+    }
+    else {
+        $body = @{
+            query   = $Query
+            options = @{
+                "`$top"  = $ProgressIncrement
+                "`$skip" = 0
+            }
+        }
     }
     if ($ScopeSplat.ManagementGroup) {
         $body.managementGroups = @($ScopeSplat.ManagementGroup)
@@ -29,13 +36,35 @@ function Search-AzGraphAllItems {
 
     [System.Collections.ArrayList] $data = [System.Collections.ArrayList]::new()
 
-    $bodyJson = $body | ConvertTo-Json -Depth 100
-    $dsi = 1
+    $dsi = 0
+
     do {
         try {
+            $bodyJson = $body | ConvertTo-Json -Depth 100
+
             $response = Invoke-AzRestMethod -Method POST `
                 -Path "/providers/Microsoft.ResourceGraph/resources?api-version=2022-10-01" `
                 -Payload $bodyJson
+            $dsi++
+                
+            if ($response.StatusCode -ge 400) {
+                $contentString = $response.Content -join ""
+                $contentObj = ConvertFrom-Json -InputObject $contentString
+                if ($contentObj.error.details[0].code -eq "ResponsePayloadTooLarge") {
+                    Write-Host "Payload too large detected. Adjusting request and retrying..."
+                    $ProgressIncrement = [Convert]::ToInt32([math]::Floor($ProgressIncrement / 2))
+                    # Modify body to include pagination options
+                    $body = @{
+                        query   = $Query
+                        options = @{
+                            "`$top"  = $ProgressIncrement
+                            "`$skip" = 0
+                        }
+                    }
+                    # Retry immediately without incrementing $dsi
+                    continue
+                }
+            }
         }
         catch {
             Write-Warning "Recovering Data Stream Error: $_"