Skip to content

Commit

Permalink
Refactored Query Planner to allow top queries.
Browse files Browse the repository at this point in the history
  • Loading branch information
michaelstaib committed Sep 6, 2023
1 parent 94af920 commit 32db9a5
Show file tree
Hide file tree
Showing 11 changed files with 209 additions and 14 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ public async Task<IQueryResult> ExecuteAsync(
operationContext.Result.SetExtension(QueryPlanHashProp, _hash);
}

// we store the context on the result for unit tests.
// we store the context on the result for unit tests.x
operationContext.Result.SetContextData(QueryPlanProp, context.QueryPlan);

// Enqueue root node to initiate the execution process.
Expand Down
9 changes: 7 additions & 2 deletions src/HotChocolate/Fusion/src/Core/Execution/Nodes/Resolve.cs
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ protected override async Task OnExecuteAsync(
{
if (state.TryGetState(SelectionSet, out var executionState))
{
var requests = new SubgraphGraphQLRequest[executionState.Count];
var requests = new SubgraphGraphQLRequest[executionState.Count];

// first we will create request for all of our selection sets.
InitializeRequests(context, executionState, requests);
Expand All @@ -57,7 +57,12 @@ protected override async Task OnExecuteAsync(
// but need to wait until the transport layer is finished and disposes the result.
context.Result.RegisterForCleanup(responses, ReturnResults);

ProcessResponses(context, executionState, requests, responses, SubgraphName);
// we need to lock the state before mutating it since there could be multiple
// query plan nodes be interested in it.
lock (executionState)
{
ProcessResponses(context, executionState, requests, responses, SubgraphName);
}
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,12 @@ protected override async Task OnExecuteAsync(
// for cleanup so that the memory can be released at the end of the execution.
context.Result.RegisterForCleanup(response, ReturnResult);

ProcessResult(context, response, batchExecutionState);
// we need to lock the state before mutating it since there could be multiple
// query plan nodes be interested in it.
lock (executionState)
{
ProcessResult(context, response, batchExecutionState);
}
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -295,10 +295,19 @@ private Subscribe CreateSubscribeNode(

private ISelectionSet ResolveSelectionSet(
QueryPlanContext context,
ExecutionStep executionStep)
=> executionStep.ParentSelection is null
SelectionExecutionStep executionStep)
{
if (executionStep.Resolver is null &&
executionStep.SelectionResolvers.Count == 0 &&
executionStep.ParentSelectionPath is not null)
{
return context.Operation.RootSelectionSet;
}

return executionStep.ParentSelection is null
? context.Operation.RootSelectionSet
: context.Operation.GetSelectionSet(
executionStep.ParentSelection,
_schema.GetType<ObjectType>(executionStep.SelectionSetTypeMetadata.Name));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ public void Invoke(QueryPlanContext context, QueryPlanDelegate next)
selectionSetType,
selections,
null,
null,
false);

while (backlog.TryDequeue(out var item))
Expand All @@ -70,6 +71,7 @@ public void Invoke(QueryPlanContext context, QueryPlanDelegate next)
item.DeclaringTypeMetadata,
item.Selections,
item.ParentSelection,
item.SelectionPath,
item.PreferBatching);
}

Expand All @@ -82,12 +84,14 @@ private void CreateExecutionSteps(
ObjectTypeMetadata selectionSetTypeMetadata,
IReadOnlyList<ISelection> selections,
ISelection? parentSelection,
SelectionPath? parentSelectionPath,
bool preferBatching)
{
var variablesInContext = new HashSet<string>();
var operation = context.Operation;
List<ISelection>? leftovers = null;

var path = new List<ISelection>();

// if this is the root selection set of a query we will
// look for some special selections.
if (!context.HasHandledSpecialQueryFields && parentSelection is null)
Expand Down Expand Up @@ -117,6 +121,7 @@ private void CreateExecutionSteps(
context.NextStepId(),
subgraph,
parentSelection,
parentSelectionPath,
_schema.GetType<IObjectType>(selectionSetTypeMetadata.Name),
selectionSetTypeMetadata);
leftovers = null;
Expand All @@ -134,6 +139,9 @@ private void CreateExecutionSteps(

foreach (var selection in current)
{
var pathIndex = path.Count;
path.Add(selection);

var field = selection.Field;
var fieldInfo = selectionSetTypeMetadata.Fields[field.Name];

Expand Down Expand Up @@ -191,10 +199,29 @@ private void CreateExecutionSteps(
backlog,
operation,
selection,
parentSelectionPath,
path,
executionStep,
preferBatching,
context.ParentSelections);
}

path.RemoveAt(pathIndex);
}

// if the current execution step has now way to resolve the data
// we will try to resolve it from the root.
if(executionStep.ParentSelection is not null &&
executionStep.ParentSelectionPath is not null &&
executionStep.Resolver is null &&
executionStep.SelectionResolvers.Count == 0)
{
if (!EnsureStepCanBeResolvedFromRoot(
executionStep.SubgraphName,
executionStep.ParentSelectionPath))
{
throw ThrowHelper.NoResolverInContext();
}
}

context.Steps.Add(executionStep);
Expand Down Expand Up @@ -258,6 +285,8 @@ private void CollectNestedSelections(
Queue<BacklogItem> backlog,
IOperation operation,
ISelection parentSelection,
SelectionPath? rootSelectionPath,
List<ISelection> path,
SelectionExecutionStep executionStep,
bool preferBatching,
Dictionary<ISelection, ISelection> parentSelectionLookup)
Expand All @@ -273,6 +302,8 @@ private void CollectNestedSelections(
backlog,
operation,
parentSelection,
rootSelectionPath,
path,
executionStep,
possibleType,
preferBatching,
Expand All @@ -284,6 +315,8 @@ private void CollectNestedSelections(
Queue<BacklogItem> backlog,
IOperation operation,
ISelection parentSelection,
SelectionPath? rootSelectionPath,
List<ISelection> path,
SelectionExecutionStep executionStep,
IObjectType possibleType,
bool preferBatching,
Expand All @@ -297,6 +330,9 @@ private void CollectNestedSelections(

foreach (var selection in selectionSet.Selections)
{
var pathIndex = path.Count;
path.Add(selection);

parentSelectionLookup.TryAdd(selection, parentSelection);
var field = declaringType.Fields[selection.Field.Name];

Expand Down Expand Up @@ -346,6 +382,8 @@ private void CollectNestedSelections(
backlog,
operation,
selection,
rootSelectionPath,
path,
executionStep,
preferBatching,
parentSelectionLookup);
Expand All @@ -355,19 +393,36 @@ private void CollectNestedSelections(
{
(leftovers ??= new()).Add(selection);
}

path.RemoveAt(pathIndex);
}

if (leftovers is not null)
{
backlog.Enqueue(
new BacklogItem(
parentSelection,
CreateSelectionPath(rootSelectionPath, path),
declaringType,
leftovers,
preferBatching));
}
}

private static SelectionPath? CreateSelectionPath(SelectionPath? rootPath, List<ISelection> pathSegments)
{
var parent = rootPath;

for (var i = 0; i < pathSegments.Count; i++)
{
parent = parent is null
? new SelectionPath(pathSegments[i])
: parent.Append(pathSegments[i]);
}

return parent;
}

private static void AddIntrospectionStepIfNotExists(
QueryPlanContext context,
IObjectField field,
Expand Down Expand Up @@ -497,6 +552,8 @@ private SelectionExecutionStep CreateNodeNestedExecutionSteps(
backlog,
operation,
nodeSelection,
null,
new List<ISelection>(),
executionStep,
entityType,
preferBatching,
Expand Down Expand Up @@ -733,14 +790,38 @@ private static bool IsNodeField(IObjectField field, IOperation operation)
field.DeclaringType.Equals(operation.RootType) &&
(field.Name.EqualsOrdinal("node") || field.Name.EqualsOrdinal("nodes"));

private bool EnsureStepCanBeResolvedFromRoot(
string subgraphName,
SelectionPath path)
{
var current = path;

while (current is not null)
{
var typeMetadata = _config.GetType<ObjectTypeMetadata>(path.Selection.DeclaringType.Name);

if (!typeMetadata.Fields[path.Selection.Field.Name].Bindings.ContainsSubgraph(subgraphName))
{
return false;
}

current = current.Parent;
}

return true;
}

private readonly struct BacklogItem(
ISelection parentSelection,
SelectionPath? selectionPath,
ObjectTypeMetadata declaringTypeMetadata,
IReadOnlyList<ISelection> selections,
bool preferBatching)
{
public ISelection ParentSelection { get; } = parentSelection;

public SelectionPath? SelectionPath { get; } = selectionPath;

public ObjectTypeMetadata DeclaringTypeMetadata { get; } = declaringTypeMetadata;

public IReadOnlyList<ISelection> Selections { get; } = selections;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -251,11 +251,21 @@ private static void TryEnqueueBatch(
private ISelectionSet ResolveSelectionSet(
QueryPlanContext context,
ExecutionStep executionStep)
=> executionStep.ParentSelection is null
{
if (executionStep is SelectionExecutionStep selectionExecStep &&
selectionExecStep.Resolver is null &&
selectionExecStep.SelectionResolvers.Count == 0 &&
selectionExecStep.ParentSelectionPath is not null)
{
return context.Operation.RootSelectionSet;
}

return executionStep.ParentSelection is null
? context.Operation.RootSelectionSet
: context.Operation.GetSelectionSet(
executionStep.ParentSelection,
_schema.GetType<Types.ObjectType>(executionStep.SelectionSetTypeMetadata.Name));
}

private readonly record struct BacklogItem(NodeAndStep[] Batch, Sequence Parent);
}
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,7 @@ private static void RegisterRequirementStep(
context.NextStepId(),
subgraph,
parentSelection,
null,
selection.DeclaringType,
typeMetadata);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,15 @@ internal RequestDocument CreateRequestDocument(
path = p;
}

if (executionStep.Resolver is null &&
executionStep.SelectionResolvers.Count == 0 &&
executionStep.ParentSelectionPath is not null)
{
rootSelectionSetNode = CreateRootLevelQuery(
executionStep.ParentSelectionPath,
rootSelectionSetNode);
}

var operationDefinitionNode = new OperationDefinitionNode(
null,
context.CreateRemoteOperationName(),
Expand All @@ -75,6 +84,26 @@ internal RequestDocument CreateRequestDocument(
path);
}

private SelectionSetNode CreateRootLevelQuery(
SelectionPath path,
SelectionSetNode selectionSet)
{
var current = path;

while (current is not null)
{
selectionSet = new SelectionSetNode(
new[]
{
current.Selection.SyntaxNode.WithSelectionSet(selectionSet)
});

current = current.Parent;
}

return selectionSet;
}

protected virtual SelectionSetNode CreateRootSelectionSetNode(
QueryPlanContext context,
SelectionExecutionStep executionStep)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ public SelectionExecutionStep(
string subgraphName,
IObjectType selectionSet,
ObjectTypeMetadata selectionSetTypeMetadata)
: this(id, subgraphName, null, selectionSet, selectionSetTypeMetadata)
: this(id, subgraphName, null, null, selectionSet, selectionSetTypeMetadata)
{
}

Expand All @@ -49,6 +49,9 @@ public SelectionExecutionStep(
/// <param name="parentSelection">
/// The parent selection of this execution step.
/// </param>
/// <param name="parentSelectionPath">
/// The selection path from which this execution step was spawned.
/// </param>
/// <param name="selectionSet">
/// The selection set that is part of this execution step.
/// </param>
Expand All @@ -59,17 +62,24 @@ public SelectionExecutionStep(
int id,
string subgraphName,
ISelection? parentSelection,
SelectionPath? parentSelectionPath,
IObjectType selectionSet,
ObjectTypeMetadata selectionSetTypeMetadata)
: base(id, parentSelection, selectionSet, selectionSetTypeMetadata)
{
SubgraphName = subgraphName;
ParentSelectionPath = parentSelectionPath;
}

/// <summary>
/// Gets the subgraph from which this execution step will fetch data.
/// </summary>
public string SubgraphName { get; }

/// <summary>
/// Gets the selection path from which this execution step was spawned.
/// </summary>
public SelectionPath? ParentSelectionPath { get; }

/// <summary>
/// Gets the resolver for this execution step.
Expand Down
Loading

0 comments on commit 32db9a5

Please sign in to comment.