diff --git a/cmd/cyphernetes/query.go b/cmd/cyphernetes/query.go index 155a2513..31ee37c6 100644 --- a/cmd/cyphernetes/query.go +++ b/cmd/cyphernetes/query.go @@ -9,6 +9,7 @@ import ( "github.com/avitaltamir/cyphernetes/pkg/core" "github.com/avitaltamir/cyphernetes/pkg/provider/apiserver" "github.com/spf13/cobra" + "gopkg.in/yaml.v3" ) var ( @@ -71,18 +72,25 @@ func runQuery(args []string, w io.Writer) { return } - // Print the results as pretty JSON. - json, err := json.MarshalIndent(results.Data, "", " ") + // Marshal data based on the output format + var output []byte + if core.OutputFormat == "yaml" { + output, err = yaml.Marshal(results.Data) + } else { // core.OutputFormat == "json" + output, err = json.MarshalIndent(results.Data, "", " ") + if !returnRawJsonOutput { + output = []byte(formatJson(string(output))) + } + } + + // Handle marshalling errors if err != nil { fmt.Fprintln(w, "Error marshalling results: ", err) return } - if !returnRawJsonOutput { - json = []byte(formatJson(string(json))) - } - if string(json) != "{}" { - fmt.Fprintln(w, string(json)) + if string(output) != "{}" { + fmt.Fprintln(w, string(output)) } } diff --git a/cmd/cyphernetes/query_test.go b/cmd/cyphernetes/query_test.go index fa4acb6e..e093f905 100644 --- a/cmd/cyphernetes/query_test.go +++ b/cmd/cyphernetes/query_test.go @@ -67,6 +67,22 @@ func TestRunQuery(t *testing.T) { } `, }, + { + name: "Successful query in YAML format", + args: []string{"MATCH (n:Pod)"}, + mockParseQuery: func(query string) (*core.Expression, error) { + return &core.Expression{}, nil + }, + mockExecute: func(expr *core.Expression, namespace string) (core.QueryResult, error) { + core.OutputFormat = "yaml" + return core.QueryResult{ + Data: map[string]interface{}{ + "test": "data", + }, + }, nil + }, + wantOut: "test: data\n\n", + }, { name: "Parse query error", args: []string{"INVALID QUERY"}, diff --git a/cmd/cyphernetes/root.go b/cmd/cyphernetes/root.go index 738eedea..8413af66 100644 --- a/cmd/cyphernetes/root.go +++ b/cmd/cyphernetes/root.go @@ -78,10 +78,21 @@ func TestExecute(args []string) error { func init() { // First set the log level rootCmd.PersistentFlags().StringVarP(&LogLevel, "loglevel", "l", "info", "The log level to use (debug, info, warn, error, fatal, panic)") - // Add a PreRun hook to set LogLevel after flag parsing - rootCmd.PersistentPreRun = func(cmd *cobra.Command, args []string) { + + // Set output format for shell and query + rootCmd.PersistentFlags().StringVar(&core.OutputFormat, "format", "json", "Output format for shell and query results") + + rootCmd.PersistentPreRunE = func(cmd *cobra.Command, args []string) error { + // Set LogLevel after flag parsing LogLevel = cmd.Flag("loglevel").Value.String() core.LogLevel = LogLevel + + // Return an error if the format is incorrect + f := cmd.Flag("format").Value.String() + if f != "yaml" && f != "json" { + return fmt.Errorf("Invalid value for --format: must be 'json' or 'yaml'") + } + return nil } rootCmd.PersistentFlags().StringVarP(&core.Namespace, "namespace", "n", "default", "The namespace to query against") diff --git a/cmd/cyphernetes/shell.go b/cmd/cyphernetes/shell.go index bfce34e8..b1dcbef0 100644 --- a/cmd/cyphernetes/shell.go +++ b/cmd/cyphernetes/shell.go @@ -19,6 +19,7 @@ import ( "github.com/avitaltamir/cyphernetes/pkg/provider/apiserver" cobra "github.com/spf13/cobra" "github.com/wader/readline" + "gopkg.in/yaml.v3" _ "k8s.io/client-go/plugin/pkg/client/auth/oidc" "k8s.io/client-go/tools/clientcmd" ) @@ -578,11 +579,24 @@ func buildDataAndGraph(resultMap map[string]interface{}, result *string, graph * } if data, ok := resultMap["Data"]; ok { - resultBytes, err := json.Marshal(data) + // Marshal data based on the output format + var output []byte + var err error + if core.OutputFormat == "yaml" { + output, err = yaml.Marshal(data) + } else { // core.OutputFormat == "json" + output, err = json.MarshalIndent(data, "", " ") + if !returnRawJsonOutput { + output = []byte(formatJson(string(output))) + } + } + + // Handle marshalling errors if err != nil { return fmt.Errorf("error marshalling data: %w", err) } - *result = string(resultBytes) + + *result = string(output) } else { *result = "{}" } diff --git a/pkg/core/k8s_query.go b/pkg/core/k8s_query.go index c9481ed8..7b70c12a 100644 --- a/pkg/core/k8s_query.go +++ b/pkg/core/k8s_query.go @@ -52,6 +52,7 @@ type QueryExecutor struct { var ( Namespace string LogLevel string + OutputFormat string AllNamespaces bool CleanOutput bool NoColor bool