address PR comments

This removes the dedicated flagSet for streaming search and move parts
of the logic to internal streaming.

Author: stefanhengl

cmd/src/search.go

@@ -17,6 +17,7 @@ import (
 
 	isatty "github.com/mattn/go-isatty"
 	"github.com/sourcegraph/src-cli/internal/api"
+	"github.com/sourcegraph/src-cli/internal/streaming"
 	"jaytaylor.com/html2text"
 )
 
@@ -50,14 +51,12 @@ Other tips:
 
 	flagSet := flag.NewFlagSet("search", flag.ExitOnError)
 	var (
-		jsonFlag        = flagSet.Bool("json", false, "Whether or not to output results as JSON")
+		jsonFlag        = flagSet.Bool("json", false, "Whether or not to output results as JSON.")
 		explainJSONFlag = flagSet.Bool("explain-json", false, "Explain the JSON output schema and exit.")
 		apiFlags        = api.NewFlags(flagSet)
-		lessFlag        = flagSet.Bool("less", true, "Pipe output to 'less -R' (only if stdout is terminal, and not json flag)")
-		streamFlag      = flagSet.Bool("stream", false, "Consume results as stream.")
-
-		// Streaming.
-		_ = flagSet.Int("display", -1, "Limit the number of results shown. Only supported for streaming.")
+		lessFlag        = flagSet.Bool("less", true, "Pipe output to 'less -R' (only if stdout is terminal, and not json flag).")
+		streamFlag      = flagSet.Bool("stream", false, "Consume results as stream. Streaming search only supports a subset of flags and parameters: trace, insecure-skip-verify, display.")
+		display         = flagSet.Int("display", -1, "Limit the number of results that are displayed. Only supported together with stream flag. Statistics continue to report all results.")
 	)
 
 	handler := func(args []string) error {
@@ -66,15 +65,12 @@ Other tips:
 		}
 
 		if *streamFlag {
-			// Remove -stream from args.
-			argsWOStream := make([]string, 0, len(args)-1)
-			for _, a := range args {
-				if a == "-stream" {
-					continue
-				}
-				argsWOStream = append(argsWOStream, a)
+			opts := streaming.Opts{
+				Display: *display,
+				Trace:   apiFlags.Trace(),
 			}
-			return streamHandler(argsWOStream)
+			client := cfg.apiClient(apiFlags, flagSet.Output())
+			return streamSearch(flagSet.Arg(0), opts, client, os.Stdout)
 		}
 
 		if *explainJSONFlag {

cmd/src/search_stream.go

@@ -2,11 +2,8 @@ package main
 
 import (
 	"bytes"
-	"context"
-	"flag"
 	"fmt"
 	"io"
-	"net/url"
 	"os"
 	"os/exec"
 	"regexp"
@@ -23,51 +20,15 @@ func init() {
 	labelRegexp, _ = regexp.Compile("(?:\\[)(.*?)(?:])")
 }
 
-// streamHandler handles search requests which contain the flag "stream".
-// Requests are sent to search/stream instead of the GraphQL api.
-func streamHandler(args []string) error {
-	flagSet := flag.NewFlagSet("streaming search", flag.ExitOnError)
-	flags := newStreamingFlags(flagSet)
-	if err := flagSet.Parse(args); err != nil {
-		return err
-	}
-
-	client := cfg.apiClient(flags.apiFlags, flagSet.Output())
-	query := flagSet.Arg(0)
-	return doStreamSearch(query, flags, client, os.Stdout)
-}
-
-func doStreamSearch(query string, flags *streamingFlags, client api.Client, w io.Writer) error {
+func streamSearch(query string, opts streaming.Opts, client api.Client, w io.Writer) error {
 	t, err := parseTemplate(streamingTemplate)
 	if err != nil {
 		panic(err)
 	}
-
-	// Create request.
-	req, err := client.NewHTTPRequest(context.Background(), "GET", "search/stream?q="+url.QueryEscape(query), nil)
-	if err != nil {
-		return err
-	}
-	req.Header.Set("Accept", "text/event-stream")
-	if flags.display >= 0 {
-		q := req.URL.Query()
-		q.Add("display", strconv.Itoa(flags.display))
-		req.URL.RawQuery = q.Encode()
-	}
-
-	// Send request.
-	resp, err := client.Do(req)
-	if err != nil {
-		return fmt.Errorf("error sending request: %w", err)
-	}
-	defer resp.Body.Close()
-
 	logError := func(msg string) {
 		_, _ = fmt.Fprintf(os.Stderr, msg)
 	}
-
-	// Process response.
-	err = streaming.Decoder{
+	decoder := streaming.Decoder{
 		OnProgress: func(progress *streaming.Progress) {
 			// We only show the final progress.
 			if !progress.Done {
@@ -157,19 +118,9 @@ func doStreamSearch(query string, flags *streamingFlags, client api.Client, w io
 				}
 			}
 		},
-	}.ReadAll(resp.Body)
-	if err != nil {
-		return fmt.Errorf("error during decoding: %w", err)
 	}
 
-	// Write trace to output.
-	if flags.Trace() {
-		_, err = fmt.Fprintf(os.Stderr, fmt.Sprintf("x-trace: %s\n", resp.Header.Get("x-trace")))
-		if err != nil {
-			return err
-		}
-	}
-	return nil
+	return streaming.Search(query, opts, client, decoder)
 }
 
 const streamingTemplate = `
@@ -412,20 +363,3 @@ func streamConvertMatchToHighlights(m streaming.EventLineMatch, isPreview bool)
 	}
 	return highlights
 }
-
-type streamingFlags struct {
-	apiFlags *api.Flags
-	display  int
-}
-
-func newStreamingFlags(flagSet *flag.FlagSet) *streamingFlags {
-	flags := &streamingFlags{
-		apiFlags: api.StreamingFlags(flagSet),
-	}
-	flagSet.IntVar(&flags.display, "display", -1, "Limit the number of results that are displayed. Note that the statistics continue to report all results.")
-	return flags
-}
-
-func (f *streamingFlags) Trace() bool {
-	return f.apiFlags.Trace()
-}

cmd/src/search_stream_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/google/go-cmp/cmp"
 
+	"github.com/sourcegraph/src-cli/internal/api"
 	"github.com/sourcegraph/src-cli/internal/streaming"
 )
 
@@ -117,11 +118,10 @@ func TestSearchStream(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	flagSet := flag.NewFlagSet("streaming search test", flag.ExitOnError)
-	flags := newStreamingFlags(flagSet)
-	client := cfg.apiClient(flags.apiFlags, flagSet.Output())
-
-	err = doStreamSearch("", flags, client, w)
+	flagSet := flag.NewFlagSet("test", flag.ExitOnError)
+	flags := api.NewFlags(flagSet)
+	client := cfg.apiClient(flags, flagSet.Output())
+	err = streamSearch("", streaming.Opts{}, client, w)
 	if err != nil {
 		t.Fatal(err)
 	}

internal/api/flags.go

@@ -38,10 +38,3 @@ func defaultFlags() *Flags {
 		insecureSkipVerify: &d,
 	}
 }
-
-func StreamingFlags(flagSet *flag.FlagSet) *Flags {
-	return &Flags{
-		trace:              flagSet.Bool("trace", false, "Log the trace ID for requests. See https://docs.sourcegraph.com/admin/observability/tracing"),
-		insecureSkipVerify: flagSet.Bool("insecure-skip-verify", false, "Skip validation of TLS certificates against trusted chains"),
-	}
-}

internal/streaming/search.go

@@ -0,0 +1,55 @@
+package streaming
+
+import (
+	"context"
+	"fmt"
+	"net/url"
+	"os"
+	"strconv"
+
+	"github.com/sourcegraph/src-cli/internal/api"
+)
+
+// Opts contains the search options supported by Search.
+type Opts struct {
+	Display int
+	Trace   bool
+}
+
+// Search calls the streaming search endpoint and uses decoder to decode the
+// response body.
+func Search(query string, opts Opts, client api.Client, decoder Decoder) error {
+	// Create request.
+	req, err := client.NewHTTPRequest(context.Background(), "GET", "search/stream?q="+url.QueryEscape(query), nil)
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Accept", "text/event-stream")
+	if opts.Display >= 0 {
+		q := req.URL.Query()
+		q.Add("display", strconv.Itoa(opts.Display))
+		req.URL.RawQuery = q.Encode()
+	}
+
+	// Send request.
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("error sending request: %w", err)
+	}
+	defer resp.Body.Close()
+
+	// Process response.
+	err = decoder.ReadAll(resp.Body)
+	if err != nil {
+		return fmt.Errorf("error during decoding: %w", err)
+	}
+
+	// Output trace.
+	if opts.Trace {
+		_, err = fmt.Fprintf(os.Stderr, fmt.Sprintf("x-trace: %s\n", resp.Header.Get("x-trace")))
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}