From 5788f0e42c4f03ce3599af74666b8e21ee9ef1d7 Mon Sep 17 00:00:00 2001
From: Satyam Singh <trueleo@protonmail.com>
Date: Thu, 19 Oct 2023 14:04:12 +0530
Subject: [PATCH] Update query command to output JSON

---
 cmd/query.go | 121 +++++++++++++++++++++++++++++++++++++++++++++++++++
 main.go      |  41 ++---------------
 2 files changed, 124 insertions(+), 38 deletions(-)
 create mode 100644 cmd/query.go

diff --git a/cmd/query.go b/cmd/query.go
new file mode 100644
index 0000000..9ec72e4
--- /dev/null
+++ b/cmd/query.go
@@ -0,0 +1,121 @@
+package cmd
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"os"
+	"pb/pkg/model"
+	"strconv"
+
+	tea "github.com/charmbracelet/bubbletea"
+	"github.com/spf13/cobra"
+)
+
+var (
+	durationFlag      = "duration"
+	durationFlagShort = "d"
+	defaultDuration   = "10"
+
+	startFlag      = "from"
+	startFlagShort = "f"
+	defaultStart   = "1m"
+
+	endFlag      = "to"
+	endFlagShort = "t"
+	defaultEnd   = "now"
+)
+
+var queryInteractive = &cobra.Command{
+	Use:     "i [stream-name] --duration 10",
+	Example: "  pb query frontend --duration 10",
+	Short:   "Interactive query table view",
+	Long:    "\n command is used to open a prompt to query a stream.",
+	Args:    cobra.ExactArgs(1),
+	PreRunE: PreRunDefaultProfile,
+	RunE: func(command *cobra.Command, args []string) error {
+		stream := args[0]
+		duration, _ := command.Flags().GetString(durationFlag)
+
+		if duration == "" {
+			duration = defaultDuration
+		}
+		durationInt, err := strconv.Atoi(duration)
+		if err != nil {
+			return err
+		}
+
+		p := tea.NewProgram(model.NewQueryModel(DefaultProfile, stream, uint(durationInt)), tea.WithAltScreen())
+		if _, err := p.Run(); err != nil {
+			fmt.Printf("there's been an error: %v", err)
+			os.Exit(1)
+		}
+
+		return nil
+	},
+}
+
+var queryJSON = &cobra.Command{
+	Use:     "query [query] --from=10m --to=now",
+	Example: "  pb query \"select * from frontend\" --from=10m --to=now",
+	Short:   "Run SQL query",
+	Long:    "\nquery command is used to run query. Output format is json string",
+	Args:    cobra.ExactArgs(1),
+	PreRunE: PreRunDefaultProfile,
+	RunE: func(command *cobra.Command, args []string) error {
+		query := args[0]
+		start, _ := command.Flags().GetString(startFlag)
+		end, _ := command.Flags().GetString(endFlag)
+
+		if start == "" {
+			start = defaultStart
+		}
+		if end == "" {
+			end = defaultEnd
+		}
+
+		client := DefaultClient()
+		return fetchData(&client, query, start, end)
+	},
+}
+
+var QueryInteractiveCmd = func() *cobra.Command {
+	queryInteractive.Flags().StringP(durationFlag, durationFlagShort, defaultDuration, "specify the duration in minutes for which queries should be executed. Defaults to 10 minutes")
+	return queryInteractive
+}()
+
+var QueryCmd = func() *cobra.Command {
+	queryJSON.Flags().StringP(startFlag, startFlagShort, defaultStart, "start time")
+	queryJSON.Flags().StringP(endFlag, endFlagShort, defaultEnd, "end time")
+	queryJSON.AddCommand(queryInteractive)
+	return queryJSON
+}()
+
+func fetchData(client *HTTPClient, query string, startTime string, endTime string) (err error) {
+	queryTemplate := `{
+    "query": "%s",
+    "startTime": "%s",
+    "endTime": "%s"
+	}
+	`
+
+	finalQuery := fmt.Sprintf(queryTemplate, query, startTime, endTime)
+
+	req, err := client.NewRequest("POST", "query", bytes.NewBuffer([]byte(finalQuery)))
+	if err != nil {
+		return
+	}
+	resp, err := client.client.Do(req)
+	if err != nil {
+		return
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		body, _ := io.ReadAll(resp.Body)
+		fmt.Println(string(body))
+	} else {
+		io.Copy(os.Stdout, resp.Body)
+	}
+	return
+}
diff --git a/main.go b/main.go
index 436e8bb..672dbaa 100644
--- a/main.go
+++ b/main.go
@@ -18,14 +18,10 @@ package main
 
 import (
 	"errors"
-	"fmt"
 	"os"
 	"pb/cmd"
 	"pb/pkg/config"
-	"pb/pkg/model"
-	"strconv"
 
-	tea "github.com/charmbracelet/bubbletea"
 	"github.com/spf13/cobra"
 )
 
@@ -36,11 +32,8 @@ var (
 )
 
 var (
-	durationFlag      = "duration"
-	durationFlagShort = "d"
-	versionFlag       = "version"
-	versionFlagShort  = "v"
-	defaultDuration   = "10"
+	versionFlag      = "version"
+	versionFlagShort = "v"
 )
 
 func defaultInitialProfile() config.Profile {
@@ -92,32 +85,6 @@ var stream = &cobra.Command{
 	PersistentPreRunE: cmd.PreRunDefaultProfile,
 }
 
-var query = &cobra.Command{
-	Use:     "query [stream-name] --duration 10",
-	Example: "  pb query frontend --duration 10",
-	Short:   "Open SQL query prompt",
-	Long:    "\nquery command is used to open a prompt to query a stream.",
-	Args:    cobra.ExactArgs(1),
-	PreRunE: cmd.PreRunDefaultProfile,
-	RunE: func(command *cobra.Command, args []string) error {
-		stream := args[0]
-		duration, _ := command.Flags().GetString(durationFlag)
-		if duration == "" {
-			duration = defaultDuration
-		}
-		durationInt, err := strconv.Atoi(duration)
-		if err != nil {
-			return err
-		}
-		p := tea.NewProgram(model.NewQueryModel(cmd.DefaultProfile, stream, uint(durationInt)), tea.WithAltScreen())
-		if _, err := p.Run(); err != nil {
-			fmt.Printf("there's been an error: %v", err)
-			os.Exit(1)
-		}
-		return nil
-	},
-}
-
 func main() {
 	profile.AddCommand(cmd.AddProfileCmd)
 	profile.AddCommand(cmd.RemoveProfileCmd)
@@ -137,10 +104,8 @@ func main() {
 	stream.AddCommand(cmd.ListStreamCmd)
 	stream.AddCommand(cmd.StatStreamCmd)
 
-	query.PersistentFlags().StringP(durationFlag, durationFlagShort, defaultDuration, "specify the duration in minutes for which queries should be executed. Defaults to 10 minutes")
-
 	cli.AddCommand(profile)
-	cli.AddCommand(query)
+	cli.AddCommand(cmd.QueryCmd)
 	cli.AddCommand(stream)
 	cli.AddCommand(user)
 	cli.AddCommand(role)