diff --git a/src/functions/private/Core/Invoke-GitHubGraphQLConnectionPrefetch.ps1 b/src/functions/private/Core/Invoke-GitHubGraphQLConnectionPrefetch.ps1 new file mode 100644 index 000000000..83ce6319e --- /dev/null +++ b/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" + } +} diff --git a/src/functions/private/Repositories/Get-GitHubMyRepositories.ps1 b/src/functions/private/Repositories/Get-GitHubMyRepositories.ps1 index 8e559a8c8..487c3a54a 100644 --- a/src/functions/private/Repositories/Get-GitHubMyRepositories.ps1 +++ b/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 { diff --git a/src/functions/private/Repositories/Get-GitHubRepositoryListByOwner.ps1 b/src/functions/private/Repositories/Get-GitHubRepositoryListByOwner.ps1 index e710278e0..fcd5c6f67 100644 --- a/src/functions/private/Repositories/Get-GitHubRepositoryListByOwner.ps1 +++ b/src/functions/private/Repositories/Get-GitHubRepositoryListByOwner.ps1 @@ -89,8 +89,6 @@ } process { - $hasNextPage = $true - $after = $null $perPageSetting = Resolve-GitHubContextSetting -Name 'PerPage' -Value $PerPage -Context $Context $graphParams = @{ PropertyList = $Property + $AdditionalProperty @@ -98,9 +96,8 @@ } $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 { diff --git a/src/functions/public/API/Invoke-GitHubGraphQLQuery.ps1 b/src/functions/public/API/Invoke-GitHubGraphQLQuery.ps1 index 3a7cd65c2..68bae74cf 100644 --- a/src/functions/public/API/Invoke-GitHubGraphQLQuery.ps1 +++ b/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