feat: view query plan as a json

router/cmd/plan_generator.go

@@ -7,6 +7,7 @@ import (
 	"log"
 	"os"
 	"os/signal"
+	"strings"
 	"syscall"
 
 	"github.com/KimMachineGun/automemlimit/memlimit"
@@ -22,7 +23,9 @@ import (
 func PlanGenerator(args []string) {
 	var planHelp bool
 
-	cfg := plan_generator.QueryPlanConfig{}
+	cfg := plan_generator.QueryPlanConfig{
+		OutputFormat: core.PlanOutputFormatText,
+	}
 	f := flag.NewFlagSet("router "+args[0], flag.ExitOnError)
 	f.BoolVar(&planHelp, "help", false, "Prints the help message")
 	f.StringVar(&cfg.ExecutionConfig, "execution-config", "", "required, execution config file location")
@@ -36,6 +39,17 @@ func PlanGenerator(args []string) {
 	f.BoolVar(&cfg.FailOnPlanError, "fail-on-error", false, "if at least one plan fails, the command exit code will be 1")
 	f.BoolVar(&cfg.FailFast, "fail-fast", false, "stop as soon as possible if a plan fails")
 	f.StringVar(&cfg.LogLevel, "log-level", "warn", "log level to use (debug, info, warn, error, panic, fatal)")
+	f.Func("print-format", "output format (text|json)", func(s string) error {
+		switch value := core.PlanOutputFormat(strings.ToLower(s)); value {
+		case "":
+			cfg.OutputFormat = core.PlanOutputFormatText
+			return nil
+		case core.PlanOutputFormatText, core.PlanOutputFormatJSON:
+			cfg.OutputFormat = value
+			return nil
+		}
+		return fmt.Errorf("must be one of: text, json (got %q)", s)
+	})
 	f.UintVar(&cfg.MaxDataSourceCollectorsConcurrency, "max-collectors", 0, "max number of concurrent data source collectors, if unset or 0, no limit will be enforced")
 
 	if err := f.Parse(args[1:]); err != nil {

router/core/plan_generator.go

@@ -2,6 +2,7 @@ package core
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"net/http"
@@ -52,6 +53,13 @@ type Planner struct {
 	operationValidator *astvalidation.OperationValidator
 }
 
+type PlanOutputFormat string
+
+const (
+	PlanOutputFormatText PlanOutputFormat = "text"
+	PlanOutputFormatJSON PlanOutputFormat = "json"
+)
+
 func NewPlanner(planConfiguration *plan.Configuration, definition *ast.Document, clientDefinition *ast.Document) (*Planner, error) {
 	planner, err := plan.NewPlanner(*planConfiguration)
 	if err != nil {
@@ -65,7 +73,7 @@ func NewPlanner(planConfiguration *plan.Configuration, definition *ast.Document,
 	}, nil
 }
 
-func (pl *Planner) PlanOperation(operationFilePath string) (string, error) {
+func (pl *Planner) PlanOperation(operationFilePath string, outputFormat PlanOutputFormat) (string, error) {
 	operation, err := pl.parseOperation(operationFilePath)
 	if err != nil {
 		return "", &PlannerOperationValidationError{err: err}
@@ -91,7 +99,18 @@ func (pl *Planner) PlanOperation(operationFilePath string) (string, error) {
 		return "", fmt.Errorf("failed to plan operation: %w", err)
 	}
 
-	return rawPlan.PrettyPrint(), nil
+	switch outputFormat {
+	case PlanOutputFormatText:
+		return rawPlan.PrettyPrint(), nil
+	case PlanOutputFormatJSON:
+		marshal, err := json.Marshal(rawPlan)
+		if err != nil {
+			return "", fmt.Errorf("failed to marshal raw plan: %w", err)
+		}
+		return string(marshal), nil
+	}
+
+	return "", fmt.Errorf("invalid type specified: %q", outputFormat)
 }
 
 func (pl *Planner) PlanParsedOperation(operation *ast.Document) (*resolve.FetchTreeQueryPlanNode, error) {

router/pkg/plan_generator/plan_generator.go

@@ -30,6 +30,7 @@ type QueryPlanConfig struct {
 	OutputReport                       bool
 	FailOnPlanError                    bool
 	FailFast                           bool
+	OutputFormat                       core.PlanOutputFormat
 	LogLevel                           string
 	Logger                             *zap.Logger
 	MaxDataSourceCollectorsConcurrency uint
@@ -52,6 +53,10 @@ func PlanGenerator(ctx context.Context, cfg QueryPlanConfig) error {
 		cfg.Concurrency = runtime.GOMAXPROCS(0)
 	}
 
+	if cfg.OutputFormat == "" {
+		cfg.OutputFormat = core.PlanOutputFormatText
+	}
+
 	queriesPath, err := filepath.Abs(cfg.SourceDir)
 	if err != nil {
 		return fmt.Errorf("failed to get absolute path for queries: %v", err)
@@ -140,7 +145,7 @@ func PlanGenerator(ctx context.Context, cfg QueryPlanConfig) error {
 
 					queryFilePath := filepath.Join(queriesPath, queryFile.Name())
 
-					outContent, err := planner.PlanOperation(queryFilePath)
+					outContent, err := planner.PlanOperation(queryFilePath, cfg.OutputFormat)
 					res := QueryPlanResult{
 						FileName: queryFile.Name(),
 						Plan:     outContent,

router/pkg/plan_generator/plan_generator_test.go

@@ -3,6 +3,7 @@ package plan_generator
 import (
 	"context"
 	"encoding/json"
+	"github.com/wundergraph/cosmo/router/core"
 	"os"
 	"path"
 	"path/filepath"
@@ -33,7 +34,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("checks queries path exists", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "notexistant"),
@@ -64,7 +67,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("checks filter file exists", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -82,7 +87,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("fail if execution config don't exists", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -99,7 +106,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("fail with invalid execution config ", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -116,7 +125,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("fails with wrong timeout duration", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -133,7 +144,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("generates a plan for every file", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -167,7 +180,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("generates a plan for every file filtered", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -195,7 +210,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("generates a result file with every plan inside", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -227,7 +244,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("will not fail on warnings and results should return the warnings and generate results file", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -260,7 +279,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("will not fail on warnings and files should have warnings and generate files", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -293,7 +314,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("when reaching timeout an error should be returned", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -310,7 +333,9 @@ func TestPlanGenerator(t *testing.T) {
 	t.Run("when reaching timeout the report should contains the error", func(t *testing.T) {
 		tempDir, err := os.MkdirTemp("", "plans-")
 		require.NoError(t, err)
-		defer os.RemoveAll(tempDir)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
 
 		cfg := QueryPlanConfig{
 			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
@@ -331,4 +356,79 @@ func TestPlanGenerator(t *testing.T) {
 		assert.NoError(t, err)
 		assert.Equal(t, errMsg, writtenResults.Error)
 	})
+
+	t.Run("generates raw json plans when Raw is enabled", func(t *testing.T) {
+		tempDir, err := os.MkdirTemp("", "plans-")
+		require.NoError(t, err)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
+
+		cfg := QueryPlanConfig{
+			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
+			OutDir:          tempDir,
+			ExecutionConfig: path.Join(getTestDataDir(), "execution_config", "base.json"),
+			Timeout:         "30s",
+			OutputFiles:     true,
+			OutputFormat:    core.PlanOutputFormatJSON,
+		}
+
+		err = PlanGenerator(context.Background(), cfg)
+		assert.NoError(t, err)
+
+		entriesInOutDir, err := os.ReadDir(tempDir)
+		assert.NoError(t, err)
+		assert.Len(t, entriesInOutDir, len(allFiles))
+
+		for _, de := range entriesInOutDir {
+			name := de.Name()
+			content, err := os.ReadFile(path.Join(tempDir, name))
+			assert.NoError(t, err)
+			var m map[string]interface{}
+
+			// One of the queries produces a failed result
+			if err := json.Unmarshal(content, &m); err != nil {
+				assert.True(t, strings.HasPrefix(string(content), "Warning:"))
+			} else {
+				assert.NotEmpty(t, m)
+			}
+		}
+	})
+
+	t.Run("generates non-raw textual plans when Raw is disabled", func(t *testing.T) {
+		tempDir, err := os.MkdirTemp("", "plans-")
+		require.NoError(t, err)
+		defer func() {
+			_ = os.RemoveAll(tempDir)
+		}()
+
+		cfg := QueryPlanConfig{
+			SourceDir:       path.Join(getTestDataDir(), "queries", "base"),
+			OutDir:          tempDir,
+			ExecutionConfig: path.Join(getTestDataDir(), "execution_config", "base.json"),
+			Timeout:         "30s",
+			OutputFiles:     true,
+			OutputFormat:    core.PlanOutputFormatText,
+		}
+
+		err = PlanGenerator(context.Background(), cfg)
+		assert.NoError(t, err)
+
+		entriesInOutDir, err := os.ReadDir(tempDir)
+		assert.NoError(t, err)
+		assert.Len(t, entriesInOutDir, len(allFiles))
+
+		for _, de := range entriesInOutDir {
+			name := de.Name()
+			assert.True(t, strings.HasSuffix(name, ".graphql"))
+			content, err := os.ReadFile(path.Join(tempDir, name))
+			assert.NoError(t, err)
+			var m map[string]interface{}
+			assert.Error(t, json.Unmarshal(content, &m))
+			// Should be textual query plan or warning
+			s := string(content)
+			assert.True(t, strings.HasPrefix(s, "QueryPlan {") || strings.HasPrefix(s, "Warning:"))
+		}
+	})
+
 }