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,
+		AllowToPlanTypeNameFields:  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,
+		AllowToPlanTypeNameFields:  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,
+		AllowToPlanTypeNameFields:  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

@@ -65,6 +65,7 @@ type fieldSelectionRewriter struct {
 	dsConfiguration    DataSource
 
 	skipFieldRefs []int
+	alwaysRewrite bool
 }
 
 type RewriteResult struct {
@@ -74,21 +75,20 @@ type RewriteResult struct {
 
 var resultNotRewritten = RewriteResult{}
 
-func newFieldSelectionRewriter(operation *ast.Document, definition *ast.Document) *fieldSelectionRewriter {
+func newFieldSelectionRewriter(operation *ast.Document, definition *ast.Document, dsConfiguration DataSource) *fieldSelectionRewriter {
+	upstreamDefinition, ok := dsConfiguration.UpstreamSchema()
+	if !ok {
+		panic("upstream schema is not defined")
+	}
 	return &fieldSelectionRewriter{
-		operation:  operation,
-		definition: definition,
+		operation:          operation,
+		definition:         definition,
+		upstreamDefinition: upstreamDefinition,
+		dsConfiguration:    dsConfiguration,
+		alwaysRewrite:      dsConfiguration.PlanningBehavior().AlwaysFlattenFragments,
 	}
 }
 
-func (r *fieldSelectionRewriter) SetUpstreamDefinition(upstreamDefinition *ast.Document) {
-	r.upstreamDefinition = upstreamDefinition
-}
-
-func (r *fieldSelectionRewriter) SetDatasourceConfiguration(dsConfiguration DataSource) {
-	r.dsConfiguration = dsConfiguration
-}
-
 func (r *fieldSelectionRewriter) RewriteFieldSelection(fieldRef int, enclosingNode ast.Node) (res RewriteResult, err error) {
 	fieldName := r.operation.FieldNameBytes(fieldRef)
 	fieldTypeNode, ok := r.definition.FieldTypeNode(fieldName, enclosingNode)
@@ -168,7 +168,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) {
@@ -346,6 +356,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 +425,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 +453,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