Run arbitraty sql

Author: ivpusic

cmd/query/helpers.go

@@ -1,8 +1,15 @@
 package query
 
 import (
+	"errors"
 	"fmt"
 	"strconv"
+	"time"
+
+	"github.com/duneanalytics/cli/output"
+	"github.com/duneanalytics/duneapi-client-go/dune"
+	"github.com/duneanalytics/duneapi-client-go/models"
+	"github.com/spf13/cobra"
 )
 
 func parseQueryID(arg string) (int, error) {
@@ -12,3 +19,40 @@ func parseQueryID(arg string) (int, error) {
 	}
 	return id, nil
 }
+
+func parsePerformance(cmd *cobra.Command) (string, error) {
+	performance, _ := cmd.Flags().GetString("performance")
+	if performance != "medium" && performance != "large" {
+		return "", fmt.Errorf("invalid performance tier %q: must be \"medium\" or \"large\"", performance)
+	}
+	return performance, nil
+}
+
+func waitAndDisplay(cmd *cobra.Command, exec dune.Execution) error {
+	resp, err := exec.WaitGetResults(5*time.Second, 60)
+	if err != nil {
+		return err
+	}
+
+	if resp.State != "QUERY_STATE_COMPLETED" {
+		msg := fmt.Sprintf("query execution failed with state %s", resp.State)
+		if resp.Error != nil {
+			msg += fmt.Sprintf(": %s", resp.Error.Message)
+		}
+		return errors.New(msg)
+	}
+
+	return output.DisplayResults(cmd, resp)
+}
+
+func displayExecuteResponse(cmd *cobra.Command, resp *models.ExecuteResponse) error {
+	w := cmd.OutOrStdout()
+	switch output.FormatFromCmd(cmd) {
+	case output.FormatJSON:
+		return output.PrintJSON(w, resp)
+	default:
+		fmt.Fprintf(w, "Execution ID: %s\n", resp.ExecutionID)
+		fmt.Fprintf(w, "State:        %s\n", resp.State)
+		return nil
+	}
+}

cmd/query/query.go

@@ -13,5 +13,6 @@ func NewQueryCmd() *cobra.Command {
 	cmd.AddCommand(newUpdateCmd())
 	cmd.AddCommand(newArchiveCmd())
 	cmd.AddCommand(newRunCmd())
+	cmd.AddCommand(newRunSQLCmd())
 	return cmd
 }

cmd/query/run.go

@@ -3,7 +3,6 @@ package query
 import (
 	"fmt"
 	"strings"
-	"time"
 
 	"github.com/duneanalytics/cli/cmdutil"
 	"github.com/duneanalytics/cli/output"
@@ -40,9 +39,9 @@ func runRun(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
-	performance, _ := cmd.Flags().GetString("performance")
-	if performance != "medium" && performance != "large" {
-		return fmt.Errorf("invalid performance tier %q: must be \"medium\" or \"large\"", performance)
+	performance, err := parsePerformance(cmd)
+	if err != nil {
+		return err
 	}
 
 	req := models.ExecuteRequest{
@@ -68,15 +67,7 @@ func runNoWait(cmd *cobra.Command, req models.ExecuteRequest) error {
 		return err
 	}
 
-	w := cmd.OutOrStdout()
-	switch output.FormatFromCmd(cmd) {
-	case output.FormatJSON:
-		return output.PrintJSON(w, resp)
-	default:
-		fmt.Fprintf(w, "Execution ID: %s\n", resp.ExecutionID)
-		fmt.Fprintf(w, "State:        %s\n", resp.State)
-		return nil
-	}
+	return displayExecuteResponse(cmd, resp)
 }
 
 func runWait(cmd *cobra.Command, req models.ExecuteRequest) error {
@@ -87,20 +78,7 @@ func runWait(cmd *cobra.Command, req models.ExecuteRequest) error {
 		return err
 	}
 
-	resp, err := exec.WaitGetResults(5*time.Second, 60)
-	if err != nil {
-		return err
-	}
-
-	if resp.State != "QUERY_STATE_COMPLETED" {
-		msg := fmt.Sprintf("query execution failed with state %s", resp.State)
-		if resp.Error != nil {
-			msg += fmt.Sprintf(": %s", resp.Error.Message)
-		}
-		return fmt.Errorf("%s", msg)
-	}
-
-	return output.DisplayResults(cmd, resp)
+	return waitAndDisplay(cmd, exec)
 }
 
 func parseParams(raw []string) (map[string]any, error) {

cmd/query/run_sql.go

@@ -0,0 +1,78 @@
+package query
+
+import (
+	"github.com/duneanalytics/cli/cmdutil"
+	"github.com/duneanalytics/cli/output"
+	"github.com/duneanalytics/duneapi-client-go/models"
+	"github.com/spf13/cobra"
+)
+
+func newRunSQLCmd() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "run-sql",
+		Short: "Execute raw SQL and display results",
+		Args:  cobra.NoArgs,
+		RunE:  runRunSQL,
+	}
+
+	cmd.Flags().String("sql", "", "SQL query to execute (required)")
+	_ = cmd.MarkFlagRequired("sql")
+	cmd.Flags().StringArray("param", nil, "query parameter in key=value format (repeatable)")
+	cmd.Flags().String("performance", "medium", `performance tier: "medium" or "large"`)
+	cmd.Flags().Int("limit", 0, "maximum number of rows to display (0 = all)")
+	cmd.Flags().Bool("no-wait", false, "submit execution and exit without waiting for results")
+	output.AddFormatFlag(cmd, "text")
+
+	return cmd
+}
+
+func runRunSQL(cmd *cobra.Command, _ []string) error {
+	sql, _ := cmd.Flags().GetString("sql")
+
+	paramFlags, _ := cmd.Flags().GetStringArray("param")
+	params, err := parseParams(paramFlags)
+	if err != nil {
+		return err
+	}
+
+	performance, err := parsePerformance(cmd)
+	if err != nil {
+		return err
+	}
+
+	req := models.ExecuteSQLRequest{
+		SQL:         sql,
+		Performance: performance,
+	}
+	if len(params) > 0 {
+		req.QueryParameters = params
+	}
+
+	noWait, _ := cmd.Flags().GetBool("no-wait")
+	if noWait {
+		return runSQLNoWait(cmd, req)
+	}
+	return runSQLWait(cmd, req)
+}
+
+func runSQLNoWait(cmd *cobra.Command, req models.ExecuteSQLRequest) error {
+	client := cmdutil.ClientFromCmd(cmd)
+
+	resp, err := client.SQLExecute(req)
+	if err != nil {
+		return err
+	}
+
+	return displayExecuteResponse(cmd, resp)
+}
+
+func runSQLWait(cmd *cobra.Command, req models.ExecuteSQLRequest) error {
+	client := cmdutil.ClientFromCmd(cmd)
+
+	exec, err := client.RunSQL(req)
+	if err != nil {
+		return err
+	}
+
+	return waitAndDisplay(cmd, exec)
+}