Use document package in vespa document commands

3 files changed, 195 insertions, 56 deletions

diff --git a/client/go/internal/cli/cmd/document.go b/client/go/internal/cli/cmd/document.go
index b5b63fd32df..07a98d2e626 100644
--- a/client/go/internal/cli/cmd/document.go
+++ b/client/go/internal/cli/cmd/document.go
@@ -5,15 +5,22 @@
 package cmd
 
 import (
+	"bytes"
+	"errors"
 	"fmt"
 	"io"
+	"net/http"
+	"os"
+	"strconv"
 	"strings"
 	"time"
 
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
+	"github.com/vespa-engine/vespa/client/go/internal/curl"
 	"github.com/vespa-engine/vespa/client/go/internal/util"
 	"github.com/vespa-engine/vespa/client/go/internal/vespa"
+	"github.com/vespa-engine/vespa/client/go/internal/vespa/document"
 )
 
 func addDocumentFlags(cmd *cobra.Command, printCurl *bool, timeoutSecs *int) {
@@ -21,6 +28,128 @@ func addDocumentFlags(cmd *cobra.Command, printCurl *bool, timeoutSecs *int) {
 	cmd.PersistentFlags().IntVarP(timeoutSecs, "timeout", "T", 60, "Timeout for the document request in seconds")
 }
 
+type serviceWithCurl struct {
+	curlCmdWriter io.Writer
+	bodyFile      string
+	service       *vespa.Service
+}
+
+func (s *serviceWithCurl) Do(request *http.Request, timeout time.Duration) (*http.Response, error) {
+	cmd, err := curl.RawArgs(request.URL.String())
+	if err != nil {
+		return nil, err
+	}
+	cmd.Method = request.Method
+	for k, vs := range request.Header {
+		for _, v := range vs {
+			cmd.Header(k, v)
+		}
+	}
+	if s.bodyFile != "" {
+		cmd.WithBodyFile(s.bodyFile)
+	}
+	cmd.Certificate = s.service.TLSOptions.CertificateFile
+	cmd.PrivateKey = s.service.TLSOptions.PrivateKeyFile
+	out := cmd.String() + "\n"
+	if _, err := io.WriteString(s.curlCmdWriter, out); err != nil {
+		return nil, err
+	}
+	return s.service.Do(request, timeout)
+}
+
+func documentClient(cli *CLI, timeoutSecs int, printCurl bool) (*document.Client, *serviceWithCurl, error) {
+	docService, err := documentService(cli)
+	if err != nil {
+		return nil, nil, err
+	}
+	service := &serviceWithCurl{curlCmdWriter: io.Discard, service: docService}
+	if printCurl {
+		service.curlCmdWriter = cli.Stderr
+	}
+	client, err := document.NewClient(document.ClientOptions{
+		Compression: document.CompressionAuto,
+		Timeout:     time.Duration(timeoutSecs) * time.Second,
+		BaseURL:     docService.BaseURL,
+		NowFunc:     time.Now,
+	}, []util.HTTPClient{service})
+	if err != nil {
+		return nil, nil, err
+	}
+	return client, service, nil
+}
+
+func sendOperation(op document.Operation, args []string, timeoutSecs int, printCurl bool, cli *CLI) error {
+	client, service, err := documentClient(cli, timeoutSecs, printCurl)
+	if err != nil {
+		return err
+	}
+	id := ""
+	filename := args[0]
+	if len(args) > 1 {
+		id = args[0]
+		filename = args[1]
+	}
+	f, err := os.Open(filename)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	doc, err := document.NewDecoder(f).Decode()
+	if errors.Is(err, document.ErrMissingId) {
+		if id == "" {
+			return fmt.Errorf("no document id given neither as argument or as a 'put', 'update' or 'remove' key in the JSON file")
+		}
+	} else if err != nil {
+		return err
+	}
+	if id != "" {
+		docId, err := document.ParseId(id)
+		if err != nil {
+			return err
+		}
+		doc.Id = docId
+	}
+	if op > -1 {
+		if id == "" && op != doc.Operation {
+			return fmt.Errorf("wanted document operation is %s, but JSON file specifies %s", op, doc.Operation)
+		}
+		doc.Operation = op
+	}
+	if doc.Body != nil {
+		service.bodyFile = f.Name()
+	}
+	result := client.Send(doc)
+	return printResult(cli, operationResult(false, doc, service.service, result), false)
+}
+
+func readDocument(id string, timeoutSecs int, printCurl bool, cli *CLI) error {
+	client, service, err := documentClient(cli, timeoutSecs, printCurl)
+	if err != nil {
+		return err
+	}
+	docId, err := document.ParseId(id)
+	if err != nil {
+		return err
+	}
+	result := client.Get(docId)
+	return printResult(cli, operationResult(true, document.Document{Id: docId}, service.service, result), true)
+}
+
+func operationResult(read bool, doc document.Document, service *vespa.Service, result document.Result) util.OperationResult {
+	bodyReader := bytes.NewReader(result.Body)
+	if result.HTTPStatus == 200 {
+		if read {
+			return util.SuccessWithPayload("Read "+doc.Id.String(), util.ReaderToJSON(bodyReader))
+		} else {
+			return util.Success(doc.Operation.String() + " " + doc.Id.String())
+		}
+	}
+	if result.HTTPStatus/100 == 4 {
+		return util.FailureWithPayload("Invalid document operation: Status "+strconv.Itoa(result.HTTPStatus), util.ReaderToJSON(bodyReader))
+	}
+	return util.FailureWithPayload(service.Description()+" at "+service.BaseURL+": Status "+strconv.Itoa(result.HTTPStatus), util.ReaderToJSON(bodyReader))
+}
+
 func newDocumentCmd(cli *CLI) *cobra.Command {
 	var (
 		printCurl   bool
@@ -44,11 +173,7 @@ should be used instead of this.`,
 		SilenceUsage:      true,
 		Args:              cobra.ExactArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
-			service, err := documentService(cli)
-			if err != nil {
-				return err
-			}
-			return printResult(cli, vespa.Send(args[0], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), false)
+			return sendOperation(-1, args, timeoutSecs, printCurl, cli)
 		},
 	}
 	addDocumentFlags(cmd, &printCurl, &timeoutSecs)
@@ -72,15 +197,7 @@ $ vespa document put id:mynamespace:music::a-head-full-of-dreams src/test/resour
 		DisableAutoGenTag: true,
 		SilenceUsage:      true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			service, err := documentService(cli)
-			if err != nil {
-				return err
-			}
-			if len(args) == 1 {
-				return printResult(cli, vespa.Put("", args[0], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), false)
-			} else {
-				return printResult(cli, vespa.Put(args[0], args[1], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), false)
-			}
+			return sendOperation(document.OperationPut, args, timeoutSecs, printCurl, cli)
 		},
 	}
 	addDocumentFlags(cmd, &printCurl, &timeoutSecs)
@@ -103,15 +220,7 @@ $ vespa document update id:mynamespace:music::a-head-full-of-dreams src/test/res
 		DisableAutoGenTag: true,
 		SilenceUsage:      true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			service, err := documentService(cli)
-			if err != nil {
-				return err
-			}
-			if len(args) == 1 {
-				return printResult(cli, vespa.Update("", args[0], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), false)
-			} else {
-				return printResult(cli, vespa.Update(args[0], args[1], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), false)
-			}
+			return sendOperation(document.OperationUpdate, args, timeoutSecs, printCurl, cli)
 		},
 	}
 	addDocumentFlags(cmd, &printCurl, &timeoutSecs)
@@ -134,14 +243,20 @@ $ vespa document remove id:mynamespace:music::a-head-full-of-dreams`,
 		DisableAutoGenTag: true,
 		SilenceUsage:      true,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			service, err := documentService(cli)
-			if err != nil {
-				return err
-			}
 			if strings.HasPrefix(args[0], "id:") {
-				return printResult(cli, vespa.RemoveId(args[0], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), false)
+				client, service, err := documentClient(cli, timeoutSecs, printCurl)
+				if err != nil {
+					return err
+				}
+				id, err := document.ParseId(args[0])
+				if err != nil {
+					return err
+				}
+				doc := document.Document{Id: id, Operation: document.OperationRemove}
+				result := client.Send(doc)
+				return printResult(cli, operationResult(false, doc, service.service, result), false)
 			} else {
-				return printResult(cli, vespa.RemoveOperation(args[0], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), false)
+				return sendOperation(document.OperationRemove, args, timeoutSecs, printCurl, cli)
 			}
 		},
 	}
@@ -162,11 +277,7 @@ func newDocumentGetCmd(cli *CLI) *cobra.Command {
 		SilenceUsage:      true,
 		Example:           `$ vespa document get id:mynamespace:music::a-head-full-of-dreams`,
 		RunE: func(cmd *cobra.Command, args []string) error {
-			service, err := documentService(cli)
-			if err != nil {
-				return err
-			}
-			return printResult(cli, vespa.Get(args[0], service, operationOptions(cli.Stderr, printCurl, timeoutSecs)), true)
+			return readDocument(args[0], timeoutSecs, printCurl, cli)
 		},
 	}
 	addDocumentFlags(cmd, &printCurl, &timeoutSecs)
@@ -181,17 +292,6 @@ func documentService(cli *CLI) (*vespa.Service, error) {
 	return cli.service(target, vespa.DocumentService, 0, cli.config.cluster())
 }
 
-func operationOptions(stderr io.Writer, printCurl bool, timeoutSecs int) vespa.OperationOptions {
-	curlOutput := io.Discard
-	if printCurl {
-		curlOutput = stderr
-	}
-	return vespa.OperationOptions{
-		CurlOutput: curlOutput,
-		Timeout:    time.Second * time.Duration(timeoutSecs),
-	}
-}
-
 func printResult(cli *CLI, result util.OperationResult, payloadOnlyOnSuccess bool) error {
 	out := cli.Stdout
 	if !result.Success {
diff --git a/client/go/internal/cli/cmd/document_test.go b/client/go/internal/cli/cmd/document_test.go
index bf9cc0404dc..00f98ee1333 100644
--- a/client/go/internal/cli/cmd/document_test.go
+++ b/client/go/internal/cli/cmd/document_test.go
@@ -5,6 +5,7 @@
 package cmd
 
 import (
+	"encoding/json"
 	"os"
 	"strconv"
 	"testing"
@@ -20,6 +21,11 @@ func TestDocumentSendPut(t *testing.T) {
 		"put", "POST", "id:mynamespace:music::a-head-full-of-dreams", "testdata/A-Head-Full-of-Dreams-Put.json", t)
 }
 
+func TestDocumentSendPutWithIdInFile(t *testing.T) {
+	assertDocumentSend([]string{"document", "testdata/A-Head-Full-of-Dreams-Put-Id.json"},
+		"put", "POST", "id:mynamespace:music::a-head-full-of-dreams", "testdata/A-Head-Full-of-Dreams-Put-Id.json", t)
+}
+
 func TestDocumentSendPutVerbose(t *testing.T) {
 	assertDocumentSend([]string{"document", "-v", "testdata/A-Head-Full-of-Dreams-Put.json"},
 		"put", "POST", "id:mynamespace:music::a-head-full-of-dreams", "testdata/A-Head-Full-of-Dreams-Put.json", t)
@@ -32,7 +38,7 @@ func TestDocumentSendUpdate(t *testing.T) {
 
 func TestDocumentSendRemove(t *testing.T) {
 	assertDocumentSend([]string{"document", "testdata/A-Head-Full-of-Dreams-Remove.json"},
-		"remove", "DELETE", "id:mynamespace:music::a-head-full-of-dreams", "testdata/A-Head-Full-of-Dreams-Remove.json", t)
+		"remove", "DELETE", "id:mynamespace:music::a-head-full-of-dreams", "", t)
 }
 
 func TestDocumentPutWithIdArg(t *testing.T) {
@@ -57,19 +63,24 @@ func TestDocumentUpdateWithoutIdArg(t *testing.T) {
 
 func TestDocumentRemoveWithIdArg(t *testing.T) {
 	assertDocumentSend([]string{"document", "remove", "id:mynamespace:music::a-head-full-of-dreams"},
-		"remove", "DELETE", "id:mynamespace:music::a-head-full-of-dreams", "testdata/A-Head-Full-of-Dreams-Remove.json", t)
+		"remove", "DELETE", "id:mynamespace:music::a-head-full-of-dreams", "", t)
 }
 
 func TestDocumentRemoveWithoutIdArg(t *testing.T) {
 	assertDocumentSend([]string{"document", "remove", "testdata/A-Head-Full-of-Dreams-Remove.json"},
-		"remove", "DELETE", "id:mynamespace:music::a-head-full-of-dreams", "testdata/A-Head-Full-of-Dreams-Remove.json", t)
+		"remove", "DELETE", "id:mynamespace:music::a-head-full-of-dreams", "", t)
+}
+
+func TestDocumentRemoveWithoutIdArgVerbose(t *testing.T) {
+	assertDocumentSend([]string{"document", "remove", "-v", "testdata/A-Head-Full-of-Dreams-Remove.json"},
+		"remove", "DELETE", "id:mynamespace:music::a-head-full-of-dreams", "", t)
 }
 
 func TestDocumentSendMissingId(t *testing.T) {
 	cli, _, stderr := newTestCLI(t)
 	assert.NotNil(t, cli.Run("document", "put", "testdata/A-Head-Full-of-Dreams-Without-Operation.json"))
 	assert.Equal(t,
-		"Error: No document id given neither as argument or as a 'put' key in the json file\n",
+		"Error: no document id given neither as argument or as a 'put', 'update' or 'remove' key in the JSON file\n",
 		stderr.String())
 }
 
@@ -77,7 +88,7 @@ func TestDocumentSendWithDisagreeingOperations(t *testing.T) {
 	cli, _, stderr := newTestCLI(t)
 	assert.NotNil(t, cli.Run("document", "update", "testdata/A-Head-Full-of-Dreams-Put.json"))
 	assert.Equal(t,
-		"Error: Wanted document operation is update but the JSON file specifies put\n",
+		"Error: wanted document operation is update, but JSON file specifies put\n",
 		stderr.String())
 }
 
@@ -103,7 +114,7 @@ func assertDocumentSend(arguments []string, expectedOperation string, expectedMe
 		t.Fatal(err)
 	}
 	expectedPath, _ := vespa.IdToURLPath(expectedDocumentId)
-	expectedURL := documentURL + "/document/v1/" + expectedPath
+	expectedURL := documentURL + "/document/v1/" + expectedPath + "?timeout=60000ms"
 
 	assert.Nil(t, cli.Run(arguments...))
 	verbose := false
@@ -113,16 +124,29 @@ func assertDocumentSend(arguments []string, expectedOperation string, expectedMe
 		}
 	}
 	if verbose {
-		expectedCurl := "curl -X " + expectedMethod + " -H 'Content-Type: application/json' --data-binary @" + expectedPayloadFile + " " + expectedURL + "\n"
+		expectedCurl := "curl -X " + expectedMethod + " -H 'Content-Type: application/json; charset=utf-8' -H 'User-Agent: Vespa CLI/0.0.0-devel'"
+		if expectedPayloadFile != "" {
+			expectedCurl += " --data-binary @" + expectedPayloadFile
+		}
+		expectedCurl += " '" + expectedURL + "'\n"
 		assert.Equal(t, expectedCurl, stderr.String())
 	}
 	assert.Equal(t, "Success: "+expectedOperation+" "+expectedDocumentId+"\n", stdout.String())
 	assert.Equal(t, expectedURL, client.LastRequest.URL.String())
-	assert.Equal(t, "application/json", client.LastRequest.Header.Get("Content-Type"))
+	assert.Equal(t, "application/json; charset=utf-8", client.LastRequest.Header.Get("Content-Type"))
 	assert.Equal(t, expectedMethod, client.LastRequest.Method)
 
-	expectedPayload, _ := os.ReadFile(expectedPayloadFile)
-	assert.Equal(t, string(expectedPayload), util.ReaderToString(client.LastRequest.Body))
+	if expectedPayloadFile != "" {
+		data, err := os.ReadFile(expectedPayloadFile)
+		assert.Nil(t, err)
+		var expectedPayload struct {
+			Fields json.RawMessage `json:"fields"`
+		}
+		assert.Nil(t, json.Unmarshal(data, &expectedPayload))
+		assert.Equal(t, `{"fields":`+string(expectedPayload.Fields)+"}", util.ReaderToString(client.LastRequest.Body))
+	} else {
+		assert.Nil(t, client.LastRequest.Body)
+	}
 }
 
 func assertDocumentGet(arguments []string, documentId string, t *testing.T) {
@@ -170,7 +194,7 @@ func assertDocumentServerError(t *testing.T, status int, errorMessage string) {
 		"id:mynamespace:music::a-head-full-of-dreams",
 		"testdata/A-Head-Full-of-Dreams-Put.json"))
 	assert.Equal(t,
-		"Error: Container (document API) at 127.0.0.1:8080: Status "+strconv.Itoa(status)+"\n\n"+errorMessage+"\n",
+		"Error: Container (document API) at http://127.0.0.1:8080: Status "+strconv.Itoa(status)+"\n\n"+errorMessage+"\n",
 		stderr.String())
 }
 
diff --git a/client/go/internal/cli/cmd/testdata/A-Head-Full-of-Dreams-Put-Id.json b/client/go/internal/cli/cmd/testdata/A-Head-Full-of-Dreams-Put-Id.json
new file mode 100644
index 00000000000..d39b22782ab
--- /dev/null
+++ b/client/go/internal/cli/cmd/testdata/A-Head-Full-of-Dreams-Put-Id.json
@@ -0,0 +1,15 @@
+{
+    "id": "id:mynamespace:music::a-head-full-of-dreams",
+    "fields": {
+        "album": "A Head Full of Dreams",
+        "artist": "Coldplay",
+        "year": 2015,
+        "category_scores": {
+            "cells": [
+                { "address" : { "cat" : "pop" },  "value": 1 },
+                { "address" : { "cat" : "rock" }, "value": 0.2 },
+                { "address" : { "cat" : "jazz" }, "value": 0 }
+            ]
+        }
+    }
+}