Merge pull request #31628 from vespa-engine/mpolden/get-field-set

Add field-set option to document get

3 files changed, 52 insertions, 36 deletions

diff --git a/client/go/internal/cli/cmd/document.go b/client/go/internal/cli/cmd/document.go
index ea99488ee90..c9f0c780be3 100644
--- a/client/go/internal/cli/cmd/document.go
+++ b/client/go/internal/cli/cmd/document.go
@@ -91,7 +91,7 @@ func sendOperation(op document.Operation, args []string, timeoutSecs int, waiter
 	return printResult(cli, operationResult(false, doc, service, result), false)
 }
 
-func readDocument(id string, timeoutSecs int, waiter *Waiter, printCurl bool, cli *CLI) error {
+func readDocument(id string, timeoutSecs int, waiter *Waiter, printCurl bool, cli *CLI, fieldSet string) error {
 	client, service, err := documentClient(cli, timeoutSecs, waiter, printCurl)
 	if err != nil {
 		return err
@@ -100,7 +100,7 @@ func readDocument(id string, timeoutSecs int, waiter *Waiter, printCurl bool, cl
 	if err != nil {
 		return err
 	}
-	result := client.Get(docId)
+	result := client.Get(docId, fieldSet)
 	return printResult(cli, operationResult(true, document.Document{Id: docId}, service, result), true)
 }
 
@@ -249,6 +249,7 @@ func newDocumentGetCmd(cli *CLI) *cobra.Command {
 		printCurl   bool
 		timeoutSecs int
 		waitSecs    int
+		fieldSet    string
 	)
 	cmd := &cobra.Command{
 		Use:               "get id",
@@ -259,9 +260,10 @@ func newDocumentGetCmd(cli *CLI) *cobra.Command {
 		Example:           `$ vespa document get id:mynamespace:music::a-head-full-of-dreams`,
 		RunE: func(cmd *cobra.Command, args []string) error {
 			waiter := cli.waiter(time.Duration(waitSecs)*time.Second, cmd)
-			return readDocument(args[0], timeoutSecs, waiter, printCurl, cli)
+			return readDocument(args[0], timeoutSecs, waiter, printCurl, cli, fieldSet)
 		},
 	}
+	cmd.Flags().StringVar(&fieldSet, "field-set", "", "Fields to include when reading document")
 	addDocumentFlags(cli, cmd, &printCurl, &timeoutSecs, &waitSecs)
 	return cmd
 }
diff --git a/client/go/internal/vespa/document/http.go b/client/go/internal/vespa/document/http.go
index 80789a208b6..25f292b92f7 100644
--- a/client/go/internal/vespa/document/http.go
+++ b/client/go/internal/vespa/document/http.go
@@ -101,42 +101,42 @@ func NewClient(options ClientOptions, httpClients []httputil.Client) (*Client, e
 	return c, nil
 }
 
-func writeQueryParam(sb *bytes.Buffer, start int, escape bool, k, v string) {
-	if sb.Len() == start {
-		sb.WriteString("?")
+func writeQueryParam(buf *bytes.Buffer, start int, escape bool, k, v string) {
+	if buf.Len() == start {
+		buf.WriteString("?")
 	} else {
-		sb.WriteString("&")
+		buf.WriteString("&")
 	}
-	sb.WriteString(k)
-	sb.WriteString("=")
+	buf.WriteString(k)
+	buf.WriteString("=")
 	if escape {
-		sb.WriteString(url.QueryEscape(v))
+		buf.WriteString(url.QueryEscape(v))
 	} else {
-		sb.WriteString(v)
+		buf.WriteString(v)
 	}
 }
 
-func (c *Client) writeDocumentPath(id Id, sb *bytes.Buffer) {
-	sb.WriteString(strings.TrimSuffix(c.options.BaseURL, "/"))
-	sb.WriteString("/document/v1/")
-	sb.WriteString(url.PathEscape(id.Namespace))
-	sb.WriteString("/")
-	sb.WriteString(url.PathEscape(id.Type))
+func (c *Client) writeDocumentPath(id Id, buf *bytes.Buffer) {
+	buf.WriteString(strings.TrimSuffix(c.options.BaseURL, "/"))
+	buf.WriteString("/document/v1/")
+	buf.WriteString(url.PathEscape(id.Namespace))
+	buf.WriteString("/")
+	buf.WriteString(url.PathEscape(id.Type))
 	if id.Number != nil {
-		sb.WriteString("/number/")
+		buf.WriteString("/number/")
 		n := uint64(*id.Number)
-		sb.WriteString(strconv.FormatUint(n, 10))
+		buf.WriteString(strconv.FormatUint(n, 10))
 	} else if id.Group != "" {
-		sb.WriteString("/group/")
-		sb.WriteString(url.PathEscape(id.Group))
+		buf.WriteString("/group/")
+		buf.WriteString(url.PathEscape(id.Group))
 	} else {
-		sb.WriteString("/docid")
+		buf.WriteString("/docid")
 	}
-	sb.WriteString("/")
-	sb.WriteString(url.PathEscape(id.UserSpecific))
+	buf.WriteString("/")
+	buf.WriteString(url.PathEscape(id.UserSpecific))
 }
 
-func (c *Client) methodAndURL(d Document, sb *bytes.Buffer) (string, string) {
+func (c *Client) methodAndURL(d Document, buf *bytes.Buffer) (string, string) {
 	httpMethod := ""
 	switch d.Operation {
 	case OperationPut:
@@ -147,28 +147,28 @@ func (c *Client) methodAndURL(d Document, sb *bytes.Buffer) (string, string) {
 		httpMethod = http.MethodDelete
 	}
 	// Base URL and path
-	c.writeDocumentPath(d.Id, sb)
+	c.writeDocumentPath(d.Id, buf)
 	// Query part
-	queryStart := sb.Len()
+	queryStart := buf.Len()
 	if c.options.Timeout > 0 {
-		writeQueryParam(sb, queryStart, false, "timeout", strconv.FormatInt(c.options.Timeout.Milliseconds(), 10)+"ms")
+		writeQueryParam(buf, queryStart, false, "timeout", strconv.FormatInt(c.options.Timeout.Milliseconds(), 10)+"ms")
 	}
 	if c.options.Route != "" {
-		writeQueryParam(sb, queryStart, true, "route", c.options.Route)
+		writeQueryParam(buf, queryStart, true, "route", c.options.Route)
 	}
 	if c.options.TraceLevel > 0 {
-		writeQueryParam(sb, queryStart, false, "tracelevel", strconv.Itoa(c.options.TraceLevel))
+		writeQueryParam(buf, queryStart, false, "tracelevel", strconv.Itoa(c.options.TraceLevel))
 	}
 	if c.options.Speedtest {
-		writeQueryParam(sb, queryStart, false, "dryRun", "true")
+		writeQueryParam(buf, queryStart, false, "dryRun", "true")
 	}
 	if d.Condition != "" {
-		writeQueryParam(sb, queryStart, true, "condition", d.Condition)
+		writeQueryParam(buf, queryStart, true, "condition", d.Condition)
 	}
 	if d.Create {
-		writeQueryParam(sb, queryStart, false, "create", "true")
+		writeQueryParam(buf, queryStart, false, "create", "true")
 	}
-	return httpMethod, sb.String()
+	return httpMethod, buf.String()
 }
 
 func (c *Client) leastBusyClient() *countingHTTPClient {
@@ -286,11 +286,14 @@ func (c *Client) Send(document Document) Result {
 }
 
 // Get retrieves document with given ID.
-func (c *Client) Get(id Id) Result {
+func (c *Client) Get(id Id, fieldSet string) Result {
 	start := c.now()
 	buf := c.buffer()
 	defer c.buffers.Put(buf)
 	c.writeDocumentPath(id, buf)
+	if fieldSet != "" {
+		writeQueryParam(buf, buf.Len(), true, "fieldSet", fieldSet)
+	}
 	url := buf.String()
 	result := Result{Id: id}
 	req, err := http.NewRequest(http.MethodGet, url, nil)
diff --git a/client/go/internal/vespa/document/http_test.go b/client/go/internal/vespa/document/http_test.go
index 89e9e96064b..74133fc73d8 100644
--- a/client/go/internal/vespa/document/http_test.go
+++ b/client/go/internal/vespa/document/http_test.go
@@ -177,7 +177,7 @@ func TestClientGet(t *testing.T) {
 }`
 	id := Id{Namespace: "mynamespace", Type: "music", UserSpecific: "doc1"}
 	httpClient.NextResponseString(200, doc)
-	result := client.Get(id)
+	result := client.Get(id, "")
 	want := Result{
 		Id:         id,
 		Body:       []byte(doc),
@@ -189,6 +189,17 @@ func TestClientGet(t *testing.T) {
 	if !reflect.DeepEqual(want, result) {
 		t.Errorf("got %+v, want %+v", result, want)
 	}
+	gotURL := httpClient.LastRequest.URL.String()
+	wantURL := "https://example.com:1337/document/v1/mynamespace/music/docid/doc1"
+	if gotURL != wantURL {
+		t.Errorf("got URL=%s, want %s", gotURL, wantURL)
+	}
+	client.Get(id, "[all]")
+	gotURL = httpClient.LastRequest.URL.String()
+	wantURL = "https://example.com:1337/document/v1/mynamespace/music/docid/doc1?fieldSet=%5Ball%5D"
+	if gotURL != wantURL {
+		t.Errorf("got URL=%s, want %s", gotURL, wantURL)
+	}
 }
 
 func TestClientSendCompressed(t *testing.T) {