feat: rewrite abstract fragments for grpc

v2/pkg/engine/datasource/graphql_datasource/graphql_datasource.go

@@ -269,14 +269,6 @@ func (p *Planner[T]) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias
 	return "", false
 }
 
-func (p *Planner[T]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior {
-	return plan.DataSourcePlanningBehavior{
-		MergeAliasedRootNodes:      true,
-		OverrideFieldPathFromAlias: true,
-		IncludeTypeNameFields:      true,
-	}
-}
-
 func (p *Planner[T]) Register(visitor *plan.Visitor, configuration plan.DataSourceConfiguration[T], dataSourcePlannerConfiguration plan.DataSourcePlannerConfiguration) error {
 
 	p.visitor = visitor
@@ -1825,6 +1817,16 @@ func (f *Factory[T]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguratio
 	return schema, true
 }
 
+func (f *Factory[T]) PlanningBehavior() plan.DataSourcePlanningBehavior {
+	b := plan.DataSourcePlanningBehavior{
+		MergeAliasedRootNodes:      true,
+		OverrideFieldPathFromAlias: true,
+		AllowPlanningTypeName:      true,
+		AlwaysFlattenFragments:     f.grpcClient != nil || f.grpcClientProvider != nil,
+	}
+	return b
+}
+
 type Source struct {
 	httpClient *http.Client
 }

v2/pkg/engine/datasource/introspection_datasource/factory.go

@@ -26,6 +26,14 @@ func (f *Factory[T]) Context() context.Context {
 	return context.TODO()
 }
 
-func (f *Factory[T]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[T]) (*ast.Document, bool) {
+func (f *Factory[T]) UpstreamSchema(_ plan.DataSourceConfiguration[T]) (*ast.Document, bool) {
 	return nil, false
 }
+
+func (f *Factory[T]) PlanningBehavior() plan.DataSourcePlanningBehavior {
+	return plan.DataSourcePlanningBehavior{
+		MergeAliasedRootNodes:      false,
+		OverrideFieldPathFromAlias: true,
+		AllowPlanningTypeName:      true,
+	}
+}

v2/pkg/engine/datasource/introspection_datasource/planner.go

@@ -45,14 +45,6 @@ func (p *Planner[T]) DownstreamResponseFieldAlias(_ int) (alias string, exists b
 	return
 }
 
-func (p *Planner[T]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior {
-	return plan.DataSourcePlanningBehavior{
-		MergeAliasedRootNodes:      false,
-		OverrideFieldPathFromAlias: true,
-		IncludeTypeNameFields:      true,
-	}
-}
-
 func (p *Planner[T]) EnterField(ref int) {
 	fieldName := p.v.Operation.FieldNameString(ref)
 	fieldAliasOrName := p.v.Operation.FieldAliasOrNameString(ref)

v2/pkg/engine/datasource/pubsub_datasource/pubsub_datasource.go

@@ -275,14 +275,6 @@ func (p *Planner[T]) ConfigureSubscription() plan.SubscriptionConfiguration {
 	return plan.SubscriptionConfiguration{}
 }
 
-func (p *Planner[T]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior {
-	return plan.DataSourcePlanningBehavior{
-		MergeAliasedRootNodes:      false,
-		OverrideFieldPathFromAlias: false,
-		IncludeTypeNameFields:      true,
-	}
-}
-
 func (p *Planner[T]) DownstreamResponseFieldAlias(_ int) (alias string, exists bool) {
 	return "", false
 }
@@ -312,10 +304,18 @@ func (f *Factory[T]) Context() context.Context {
 	return f.executionContext
 }
 
-func (f *Factory[T]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[T]) (*ast.Document, bool) {
+func (f *Factory[T]) UpstreamSchema(_ plan.DataSourceConfiguration[T]) (*ast.Document, bool) {
 	return nil, false
 }
 
+func (f *Factory[T]) PlanningBehavior() plan.DataSourcePlanningBehavior {
+	return plan.DataSourcePlanningBehavior{
+		MergeAliasedRootNodes:      false,
+		OverrideFieldPathFromAlias: false,
+		AllowPlanningTypeName:      true,
+	}
+}
+
 func buildEventDataBytes(ref int, visitor *plan.Visitor, variables *resolve.Variables) ([]byte, error) {
 	// Collect the field arguments for fetch based operations
 	fieldArgs := visitor.Operation.FieldArguments(ref)

v2/pkg/engine/datasource/staticdatasource/static_datasource.go

@@ -12,24 +12,28 @@ import (
 	"github.com/wundergraph/graphql-go-tools/v2/pkg/engine/resolve"
 )
 
-type Configuration struct {
-	Data string `json:"data"`
-}
-
 type Factory[T Configuration] struct{}
 
-func (f *Factory[T]) Planner(logger abstractlogger.Logger) plan.DataSourcePlanner[T] {
+func (f *Factory[T]) Planner(_ abstractlogger.Logger) plan.DataSourcePlanner[T] {
 	return &Planner[T]{}
 }
 
 func (f *Factory[T]) Context() context.Context {
 	return context.TODO()
 }
 
-func (f *Factory[T]) UpstreamSchema(dataSourceConfig plan.DataSourceConfiguration[T]) (*ast.Document, bool) {
+func (f *Factory[T]) UpstreamSchema(_ plan.DataSourceConfiguration[T]) (*ast.Document, bool) {
 	return nil, false
 }
 
+func (f *Factory[T]) PlanningBehavior() plan.DataSourcePlanningBehavior {
+	return plan.DataSourcePlanningBehavior{}
+}
+
+type Configuration struct {
+	Data string `json:"data"`
+}
+
 type Planner[T Configuration] struct {
 	id     int
 	config Configuration
@@ -47,14 +51,6 @@ func (p *Planner[T]) DownstreamResponseFieldAlias(downstreamFieldRef int) (alias
 	// skip, not required
 	return
 }
-
-func (p *Planner[T]) DataSourcePlanningBehavior() plan.DataSourcePlanningBehavior {
-	return plan.DataSourcePlanningBehavior{
-		MergeAliasedRootNodes:      false,
-		OverrideFieldPathFromAlias: false,
-	}
-}
-
 func (p *Planner[T]) Register(_ *plan.Visitor, configuration plan.DataSourceConfiguration[T], _ plan.DataSourcePlannerConfiguration) error {
 	p.config = Configuration(configuration.CustomConfiguration())
 	return nil

v2/pkg/engine/plan/abstract_selection_rewriter.go

@@ -13,9 +13,11 @@ import (
 )
 
 var (
-	FieldDoesntHaveSelectionSetErr          = errors.New("unexpected error: field does not have a selection set")
-	InlineFragmentDoesntHaveSelectionSetErr = errors.New("unexpected error: inline fragment does not have a selection set")
-	InlineFragmentTypeIsNotExistsErr        = errors.New("unexpected error: inline fragment type condition does not exists")
+	ErrFieldHasNoSelectionSet          = errors.New("unexpected error: field does not have a selection set")
+	ErrInlineFragmentHasNoSelectionSet = errors.New("unexpected error: inline fragment does not have a selection set")
+	ErrInlineFragmentHasNoCondition    = errors.New("unexpected error: inline fragment type condition does not exist")
+
+	ErrNoUpstreamSchema = errors.New("unexpected error: upstream schema is not defined in DataSource")
 )
 
 /*
@@ -65,6 +67,7 @@ type fieldSelectionRewriter struct {
 	dsConfiguration    DataSource
 
 	skipFieldRefs []int
+	alwaysRewrite bool
 }
 
 type RewriteResult struct {
@@ -74,19 +77,18 @@ type RewriteResult struct {
 
 var resultNotRewritten = RewriteResult{}
 
-func newFieldSelectionRewriter(operation *ast.Document, definition *ast.Document) *fieldSelectionRewriter {
-	return &fieldSelectionRewriter{
-		operation:  operation,
-		definition: definition,
+func newFieldSelectionRewriter(operation *ast.Document, definition *ast.Document, dsConfiguration DataSource) (*fieldSelectionRewriter, error) {
+	upstreamDefinition, ok := dsConfiguration.UpstreamSchema()
+	if !ok {
+		return nil, ErrNoUpstreamSchema
 	}
-}
-
-func (r *fieldSelectionRewriter) SetUpstreamDefinition(upstreamDefinition *ast.Document) {
-	r.upstreamDefinition = upstreamDefinition
-}
-
-func (r *fieldSelectionRewriter) SetDatasourceConfiguration(dsConfiguration DataSource) {
-	r.dsConfiguration = dsConfiguration
+	return &fieldSelectionRewriter{
+		operation:          operation,
+		definition:         definition,
+		upstreamDefinition: upstreamDefinition,
+		dsConfiguration:    dsConfiguration,
+		alwaysRewrite:      dsConfiguration.PlanningBehavior().AlwaysFlattenFragments,
+	}, nil
 }
 
 func (r *fieldSelectionRewriter) RewriteFieldSelection(fieldRef int, enclosingNode ast.Node) (res RewriteResult, err error) {
@@ -168,7 +170,17 @@ func (r *fieldSelectionRewriter) processUnionSelection(fieldRef int, unionDefRef
 	}, nil
 }
 
+func (r *fieldSelectionRewriter) mustRewrite(s selectionSetInfo) bool {
+	return r.alwaysRewrite &&
+		(s.hasInlineFragmentsOnInterfaces ||
+			s.hasInlineFragmentsOnUnions ||
+			s.hasInlineFragmentsOnObjects)
+}
+
 func (r *fieldSelectionRewriter) unionFieldSelectionNeedsRewrite(selectionSetInfo selectionSetInfo, unionTypeNames, entityNames []string) (needRewrite bool) {
+	if r.mustRewrite(selectionSetInfo) {
+		return true
+	}
 	if selectionSetInfo.hasInlineFragmentsOnObjects {
 		// when we have types not exists in the current datasource - we need to rewrite
 		if r.objectFragmentsRequiresCleanup(selectionSetInfo.inlineFragmentsOnObjects, unionTypeNames) {
@@ -263,7 +275,7 @@ func (r *fieldSelectionRewriter) replaceFieldSelections(fieldRef int, newSelecti
 func (r *fieldSelectionRewriter) processObjectSelection(fieldRef int, objectDefRef int) (res RewriteResult, err error) {
 	selectionSetRef, ok := r.operation.FieldSelectionSet(fieldRef)
 	if !ok {
-		return resultNotRewritten, FieldDoesntHaveSelectionSetErr
+		return resultNotRewritten, ErrFieldHasNoSelectionSet
 	}
 
 	fieldTypeName := r.definition.ObjectTypeDefinitionNameBytes(objectDefRef)
@@ -346,6 +358,9 @@ func (r *fieldSelectionRewriter) rewriteObjectSelection(fieldRef int, fieldInfo
 }
 
 func (r *fieldSelectionRewriter) objectFieldSelectionNeedsRewrite(selectionSetInfo selectionSetInfo, objectTypeName string) (needRewrite bool) {
+	if r.mustRewrite(selectionSetInfo) {
+		return true
+	}
 	if selectionSetInfo.hasInlineFragmentsOnObjects {
 		if r.objectFragmentsRequiresCleanup(selectionSetInfo.inlineFragmentsOnObjects, []string{objectTypeName}) {
 			return true
@@ -412,6 +427,10 @@ func (r *fieldSelectionRewriter) processInterfaceSelection(fieldRef int, interfa
 }
 
 func (r *fieldSelectionRewriter) interfaceFieldSelectionNeedsRewrite(selectionSetInfo selectionSetInfo, interfaceTypeNames []string, entityNames []string) (needRewrite bool) {
+	if r.mustRewrite(selectionSetInfo) {
+		return true
+	}
+
 	// when we do not have fragments
 	if !selectionSetInfo.hasInlineFragmentsOnInterfaces &&
 		!selectionSetInfo.hasInlineFragmentsOnUnions &&
@@ -436,9 +455,10 @@ func (r *fieldSelectionRewriter) interfaceFieldSelectionNeedsRewrite(selectionSe
 			return true
 		}
 
-		// in case it is an interface object, and we have fragments on concrete types - we have to add shared __typename selection
-		// it will mean that we will rewrite a query to separate concrete type fragments, but due to nature of the interface object
-		// they eventually will be flattened by datasource into a single fragment or just a flatten query.
+		// In case it is an interface object, and we have fragments on concrete types - we have to add shared __typename selection.
+		// It will mean that we will rewrite a query to separate concrete type fragments,
+		// but due to the nature of the interface object, they eventually will be flattened by datasource
+		// into a single fragment or just a flattened query.
 		// So it should be safe to rewrite a field.
 		if selectionSetInfo.isInterfaceObject {
 			return !selectionSetInfo.hasTypeNameSelection
@@ -618,7 +638,7 @@ func (r *fieldSelectionRewriter) collectChangedRefs(fieldRef int, fieldRefsPaths
 	for fieldRef, path := range fieldRefsPaths {
 		newRefs, ok := pathsToRefs[path]
 		if !ok {
-			// TODO: some path actually could dissapear due to rewrite
+			// TODO: some paths could actually disappear due to rewrite
 			continue
 		}
 

v2/pkg/engine/plan/abstract_selection_rewriter_info.go

@@ -83,7 +83,7 @@ func (r *fieldSelectionRewriter) selectionSetFieldSelections(selectionSetRef int
 func (r *fieldSelectionRewriter) collectFieldInformation(fieldRef int) (selectionSetInfo, error) {
 	fieldSelectionSetRef, ok := r.operation.FieldSelectionSet(fieldRef)
 	if !ok {
-		return selectionSetInfo{}, FieldDoesntHaveSelectionSetErr
+		return selectionSetInfo{}, ErrFieldHasNoSelectionSet
 	}
 
 	return r.collectSelectionSetInformation(fieldSelectionSetRef)
@@ -101,7 +101,7 @@ func (r *fieldSelectionRewriter) collectInlineFragmentInformation(
 	typeCondition := r.operation.InlineFragmentTypeConditionNameString(inlineFragmentRef)
 	inlineFragmentSelectionSetRef, ok := r.operation.InlineFragmentSelectionSet(inlineFragmentRef)
 	if !ok {
-		return InlineFragmentDoesntHaveSelectionSetErr
+		return ErrInlineFragmentHasNoSelectionSet
 	}
 
 	hasDirectives := r.operation.InlineFragmentHasDirectives(inlineFragmentRef)
@@ -110,7 +110,7 @@ func (r *fieldSelectionRewriter) collectInlineFragmentInformation(
 	// because it could be absent in the current SUBGRAPH document
 	definitionNode, hasNode := r.definition.NodeByNameStr(typeCondition)
 	if !hasNode {
-		return InlineFragmentTypeIsNotExistsErr
+		return ErrInlineFragmentHasNoCondition
 	}
 
 	selectionSetInfo, err := r.collectSelectionSetInformation(inlineFragmentSelectionSetRef)
@@ -140,7 +140,7 @@ func (r *fieldSelectionRewriter) collectInlineFragmentInformation(
 			typeNamesImplementingInterface: typeNamesImplementingInterface,
 		}
 
-		// NOTE: We are getting type names implementing interface from the current SUBGRAPH definion
+		// NOTE: We are getting type names implementing interface from the current SUBGRAPH definition
 		// NOTE: at this point we ignore case when upstreamNode is not exists in the upstream schema
 		upstreamNode, hasUpstreamNode := r.upstreamDefinition.NodeByNameStr(typeCondition)
 		if hasUpstreamNode {
@@ -163,7 +163,7 @@ func (r *fieldSelectionRewriter) collectInlineFragmentInformation(
 			unionMemberTypeNames:    unionMemberTypeNames,
 		}
 
-		// NOTE: We are getting type names of union members from the current SUBGRAPH definion
+		// NOTE: We are getting type names of union members from the current SUBGRAPH definition
 		// NOTE: at this point we ignore case when upstreamNode is not exists in the upstream schema
 		upstreamNode, hasUpstreamNode := r.upstreamDefinition.NodeByNameStr(typeCondition)
 		if hasUpstreamNode {