Prototype GraphQL connection prefetch

src/functions/private/Core/Invoke-GitHubGraphQLConnectionPrefetch.ps1

@@ -0,0 +1,179 @@
+function Invoke-GitHubGraphQLConnectionPrefetch {
+    <#
+        .SYNOPSIS
+        Prefetches paginated GraphQL connection nodes in a background runspace.
+
+        .DESCRIPTION
+        Uses a background runspace and a BlockingCollection queue to fetch pages from a GraphQL connection
+        while the foreground pipeline emits previously fetched nodes.
+    #>
+    [CmdletBinding()]
+    param(
+        [Parameter(Mandatory)]
+        [string] $Query,
+
+        [Parameter()]
+        [hashtable] $Variables = @{},
+
+        [Parameter(Mandatory)]
+        [string[]] $ConnectionPath,
+
+        [Parameter(Mandatory)]
+        [object] $Context,
+
+        [Parameter()]
+        [int] $QueueCapacity = 500
+    )
+
+    begin {
+        $stackPath = Get-PSCallStackPath
+        Write-Debug "[$stackPath] - Start"
+    }
+
+    process {
+        Update-GitHubUserAccessToken -Context $Context | Out-Null
+
+        $apiBaseUri = Resolve-GitHubContextSetting -Name 'ApiBaseUri' -Context $Context
+        $apiVersion = Resolve-GitHubContextSetting -Name 'ApiVersion' -Context $Context
+        $endpoint = New-Uri -BaseUri $apiBaseUri -Path '/graphql' -AsString
+        $token = $Context.Token | ConvertFrom-SecureString -AsPlainText
+
+        $workerVariables = @{}
+        foreach ($key in $Variables.Keys) {
+            $workerVariables[$key] = $Variables[$key]
+        }
+
+        $queue = [System.Collections.Concurrent.BlockingCollection[object]]::new($QueueCapacity)
+        $workerErrorQueue = [System.Collections.Concurrent.ConcurrentQueue[object]]::new()
+        $cancellation = [System.Threading.CancellationTokenSource]::new()
+
+        $workerRunspace = [runspacefactory]::CreateRunspace()
+        $workerRunspace.Open()
+
+        $workerPowerShell = [powershell]::Create()
+        $workerPowerShell.Runspace = $workerRunspace
+
+        $producerScript = {
+            param(
+                [System.Collections.Concurrent.BlockingCollection[object]] $Queue,
+                [System.Collections.Concurrent.ConcurrentQueue[object]] $WorkerErrorQueue,
+                [System.Threading.CancellationToken] $CancellationToken,
+                [string] $Endpoint,
+                [string] $ApiVersion,
+                [string] $UserAgent,
+                [string] $Token,
+                [string] $Query,
+                [hashtable] $Variables,
+                [string[]] $ConnectionPath
+            )
+
+            try {
+                $headers = @{
+                    Accept                 = 'application/vnd.github+json; charset=utf-8'
+                    'X-GitHub-Api-Version' = $ApiVersion
+                    'User-Agent'           = $UserAgent
+                    Authorization          = "Bearer $Token"
+                }
+
+                $hasNextPage = $true
+                while ($hasNextPage -and -not $CancellationToken.IsCancellationRequested) {
+                    $body = @{
+                        query     = $Query
+                        variables = $Variables
+                    } | ConvertTo-Json -Depth 100
+
+                    $response = Invoke-RestMethod -Uri $Endpoint -Method Post -Headers $headers -ContentType 'application/vnd.github+json; charset=utf-8' -Body $body -ErrorAction Stop
+
+                    if ($response.errors) {
+                        throw "GraphQL prefetch worker failed: $($response.errors | ConvertTo-Json -Depth 20 -Compress)"
+                    }
+
+                    $connection = $response.data
+                    foreach ($segment in $ConnectionPath) {
+                        if ($null -eq $connection) {
+                            break
+                        }
+
+                        $connection = $connection.$segment
+                    }
+
+                    if ($null -eq $connection) {
+                        throw "GraphQL prefetch worker could not resolve connection path: $($ConnectionPath -join '.')"
+                    }
+
+                    if ($null -ne $connection.nodes) {
+                        foreach ($node in $connection.nodes) {
+                            $Queue.Add($node, $CancellationToken)
+                        }
+                    }
+
+                    $hasNextPage = [bool]$connection.pageInfo.hasNextPage
+                    $Variables['Cursor'] = $connection.pageInfo.endCursor
+                }
+            } catch {
+                $WorkerErrorQueue.Enqueue($_)
+            } finally {
+                $Queue.CompleteAdding()
+            }
+        }
+
+        $null = $workerPowerShell.AddScript($producerScript.ToString()).
+            AddArgument($queue).
+            AddArgument($workerErrorQueue).
+            AddArgument($cancellation.Token).
+            AddArgument($endpoint).
+            AddArgument($apiVersion).
+            AddArgument($script:UserAgent).
+            AddArgument($token).
+            AddArgument($Query).
+            AddArgument($workerVariables).
+            AddArgument($ConnectionPath)
+
+        $workerAsyncResult = $workerPowerShell.BeginInvoke()
+        $completedNormally = $false
+
+        try {
+            foreach ($item in $queue.GetConsumingEnumerable()) {
+                Write-Output $item
+            }
+
+            $workerPowerShell.EndInvoke($workerAsyncResult)
+
+            $workerError = $null
+            if ($workerErrorQueue.TryDequeue([ref]$workerError)) {
+                throw $workerError
+            }
+
+            $completedNormally = $true
+        } finally {
+            if (-not $completedNormally) {
+                $cancellation.Cancel()
+
+                if (-not $queue.IsAddingCompleted) {
+                    $queue.CompleteAdding()
+                }
+
+                try {
+                    $workerPowerShell.Stop()
+                } catch {
+                }
+
+                try {
+                    if ($workerAsyncResult) {
+                        $workerPowerShell.EndInvoke($workerAsyncResult)
+                    }
+                } catch {
+                }
+            }
+
+            $workerPowerShell.Dispose()
+            $workerRunspace.Dispose()
+            $cancellation.Dispose()
+            $queue.Dispose()
+        }
+    }
+
+    end {
+        Write-Debug "[$stackPath] - End"
+    }
+}

src/functions/private/Repositories/Get-GitHubMyRepositories.ps1

@@ -78,8 +78,6 @@
     }
 
     process {
-        $hasNextPage = $true
-        $after = $null
         $perPageSetting = Resolve-GitHubContextSetting -Name 'PerPage' -Value $PerPage -Context $Context
 
         # CustomProperties are only available for organization-owned repos; viewer repos are user-owned.
@@ -90,9 +88,8 @@
         }
         $graphQLFields = ConvertTo-GitHubGraphQLField @graphParams
 
-        do {
-            $apiParams = @{
-                Query     = @"
+        $apiParams = @{
+            Query          = @"
 query(
     `$PerPage: Int!,
     `$Cursor: String,
@@ -121,25 +118,21 @@ $graphQLFields
   }
 }
 "@
-                Variables = @{
-                    PerPage      = $perPageSetting
-                    Cursor       = $after
-                    Affiliations = $Affiliation | ForEach-Object { $_.ToString().ToUpper() }
-                    Visibility   = -not [string]::IsNullOrEmpty($Visibility) ? $Visibility.ToString().ToUpper() : $null
-                    IsArchived   = $IsArchived
-                    IsFork       = $IsFork
-                }
-                Context   = $Context
+            Variables      = @{
+                PerPage      = $perPageSetting
+                Cursor       = $null
+                Affiliations = $Affiliation | ForEach-Object { $_.ToString().ToUpper() }
+                Visibility   = -not [string]::IsNullOrEmpty($Visibility) ? $Visibility.ToString().ToUpper() : $null
+                IsArchived   = $IsArchived
+                IsFork       = $IsFork
             }
+            ConnectionPath = @('viewer', 'repositories')
+            Context        = $Context
+        }
 
-            Invoke-GitHubGraphQLQuery @apiParams | ForEach-Object {
-                $_.viewer.repositories.nodes | ForEach-Object {
-                    [GitHubRepository]::new($_)
-                }
-                $hasNextPage = $response.pageInfo.hasNextPage
-                $after = $response.pageInfo.endCursor
-            }
-        } while ($hasNextPage)
+        Invoke-GitHubGraphQLQuery @apiParams | ForEach-Object {
+            [GitHubRepository]::new($_)
+        }
     }
 
     end {

src/functions/private/Repositories/Get-GitHubRepositoryListByOwner.ps1

@@ -89,18 +89,15 @@
     }
 
     process {
-        $hasNextPage = $true
-        $after = $null
         $perPageSetting = Resolve-GitHubContextSetting -Name 'PerPage' -Value $PerPage -Context $Context
         $graphParams = @{
             PropertyList         = $Property + $AdditionalProperty
             PropertyToGraphQLMap = [GitHubRepository]::PropertyToGraphQLMap
         }
         $graphQLFields = ConvertTo-GitHubGraphQLField @graphParams
 
-        do {
-            $apiParams = @{
-                Query     = @"
+        $apiParams = @{
+            Query          = @"
 query(
     `$Owner: String!,
     `$PerPage: Int!,
@@ -134,27 +131,23 @@ query(
   }
 }
 "@
-                Variables = @{
-                    Owner             = $Owner
-                    PerPage           = $perPageSetting
-                    Cursor            = $after
-                    Affiliations      = [string]::IsNullOrEmpty($Affiliation) ? $null : $Affiliation.ToUpper()
-                    OwnerAffiliations = [string]::IsNullOrEmpty($OwnerAffiliations) ? $null : $OwnerAffiliations.ToUpper()
-                    Visibility        = [string]::IsNullOrEmpty($Visibility) ? $null : $Visibility.ToUpper()
-                    IsArchived        = $IsArchived
-                    IsFork            = $IsFork
-                }
-                Context   = $Context
+            Variables      = @{
+                Owner             = $Owner
+                PerPage           = $perPageSetting
+                Cursor            = $null
+                Affiliations      = [string]::IsNullOrEmpty($Affiliation) ? $null : $Affiliation.ToUpper()
+                OwnerAffiliations = [string]::IsNullOrEmpty($OwnerAffiliations) ? $null : $OwnerAffiliations.ToUpper()
+                Visibility        = [string]::IsNullOrEmpty($Visibility) ? $null : $Visibility.ToUpper()
+                IsArchived        = $IsArchived
+                IsFork            = $IsFork
             }
+            ConnectionPath = @('repositoryOwner', 'repositories')
+            Context        = $Context
+        }
 
-            Invoke-GitHubGraphQLQuery @apiParams | ForEach-Object {
-                foreach ($repository in $_.repositoryOwner.repositories.nodes) {
-                    [GitHubRepository]::new($repository)
-                }
-                $hasNextPage = $_.repositoryOwner.repositories.pageInfo.hasNextPage
-                $after = $_.repositoryOwner.repositories.pageInfo.endCursor
-            }
-        } while ($hasNextPage)
+        Invoke-GitHubGraphQLQuery @apiParams | ForEach-Object {
+            [GitHubRepository]::new($_)
+        }
     }
 
     end {

src/functions/public/API/Invoke-GitHubGraphQLQuery.ps1

@@ -27,6 +27,14 @@
         [Parameter()]
         [hashtable] $Variables,
 
+        # Connection path to stream as prefetched nodes instead of returning the raw GraphQL data object.
+        [Parameter()]
+        [string[]] $ConnectionPath,
+
+        # Maximum number of prefetched nodes to hold in memory before applying backpressure to the producer.
+        [Parameter()]
+        [int] $QueueCapacity = 500,
+
         # The context to run the command in. Used to get the details for the API call.
         # Can be either a string or a GitHubContext object.
         [Parameter()]
@@ -41,6 +49,11 @@
     }
 
     process {
+        if ($ConnectionPath) {
+            Invoke-GitHubGraphQLConnectionPrefetch -Query $Query -Variables $Variables -ConnectionPath $ConnectionPath -QueueCapacity $QueueCapacity -Context $Context
+            return
+        }
+
         $body = @{
             query     = $Query
             variables = $Variables