add "Get" method for es client (#81)

Author: mrparkers

go/v1beta1/storage/esutil/client.go

@@ -25,6 +25,8 @@ import (
 	"go.uber.org/zap"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
+	"net/http"
+	"net/url"
 )
 
 type CreateRequest struct {
@@ -50,6 +52,11 @@ type MultiSearchRequest struct {
 	Searches []*EsSearch
 }
 
+type GetRequest struct {
+	Index      string
+	DocumentId string
+}
+
 type MultiGetRequest struct {
 	Index       string
 	DocumentIds []string
@@ -93,6 +100,7 @@ type Client interface {
 	BulkCreate(ctx context.Context, request *BulkCreateRequest) (*EsBulkResponse, error)
 	Search(ctx context.Context, request *SearchRequest) (*SearchResponse, error)
 	MultiSearch(ctx context.Context, request *MultiSearchRequest) (*EsMultiSearchResponse, error)
+	Get(ctx context.Context, request *GetRequest) (*EsGetResponse, error)
 	MultiGet(ctx context.Context, request *MultiGetRequest) (*EsMultiGetResponse, error)
 	Update(ctx context.Context, request *UpdateRequest) error
 	Delete(ctx context.Context, request *DeleteRequest) error
@@ -350,6 +358,31 @@ func (c *client) MultiSearch(ctx context.Context, request *MultiSearchRequest) (
 	return &response, nil
 }
 
+func (c *client) Get(ctx context.Context, request *GetRequest) (*EsGetResponse, error) {
+	log := c.logger.Named("Get").With(zap.String("index", request.Index), zap.String("documentId", request.DocumentId))
+
+	res, err := c.esClient.Get(
+		request.Index,
+		url.QueryEscape(request.DocumentId),
+		c.esClient.Get.WithContext(ctx),
+	)
+	if err != nil {
+		return nil, err
+	}
+	if res.IsError() && res.StatusCode != http.StatusNotFound {
+		return nil, fmt.Errorf("unexpected response from elasticsearch: %s", res.String())
+	}
+
+	var response EsGetResponse
+	if err = DecodeResponse(res.Body, &response); err != nil {
+		return nil, err
+	}
+
+	log.Debug("elasticsearch response", zap.Any("response", response))
+
+	return &response, nil
+}
+
 func (c *client) MultiGet(ctx context.Context, request *MultiGetRequest) (*EsMultiGetResponse, error) {
 	log := c.logger.Named("MultiGet")
 

go/v1beta1/storage/esutil/client_test.go

@@ -33,6 +33,7 @@ import (
 	"io"
 	"io/ioutil"
 	"net/http"
+	"net/url"
 	"strconv"
 	"strings"
 )
@@ -605,6 +606,97 @@ var _ = Describe("elasticsearch client", func() {
 		})
 	})
 
+	Context("Get", func() {
+		var (
+			expectedDocumentId string
+			expectedIndex      string
+
+			expectedGetRequest  *GetRequest
+			expectedGetResponse *EsGetResponse
+
+			actualGetResponse *EsGetResponse
+			actualErr         error
+		)
+
+		BeforeEach(func() {
+			expectedDocumentId = fake.LetterN(10)
+			expectedIndex = fake.LetterN(10)
+
+			expectedGetRequest = &GetRequest{
+				Index:      expectedIndex,
+				DocumentId: expectedDocumentId,
+			}
+
+			expectedOccurrence := createRandomOccurrence()
+			expectedOccurrenceJson, _ := protojson.Marshal(protov1.MessageV2(expectedOccurrence))
+
+			expectedGetResponse = &EsGetResponse{
+				Id:     expectedDocumentId,
+				Found:  true,
+				Source: expectedOccurrenceJson,
+			}
+
+			transport.PreparedHttpResponses = []*http.Response{
+				{
+					StatusCode: http.StatusOK,
+					Body:       structToJsonBody(expectedGetResponse),
+				},
+			}
+		})
+
+		JustBeforeEach(func() {
+			actualGetResponse, actualErr = client.Get(ctx, expectedGetRequest)
+		})
+
+		It("should send the get request to ES", func() {
+			Expect(transport.ReceivedHttpRequests[0].Method).To(Equal(http.MethodGet))
+			Expect(transport.ReceivedHttpRequests[0].URL.Path).To(Equal(fmt.Sprintf("/%s/_doc/%s", expectedIndex, expectedDocumentId)))
+		})
+
+		It("should return the response and no error", func() {
+			Expect(actualErr).ToNot(HaveOccurred())
+			Expect(actualGetResponse).To(Equal(expectedGetResponse))
+		})
+
+		When("the get operation fails", func() {
+			BeforeEach(func() {
+				transport.PreparedHttpResponses = []*http.Response{
+					{
+						StatusCode: http.StatusInternalServerError,
+					},
+				}
+			})
+
+			It("should return an error", func() {
+				Expect(actualErr).To(HaveOccurred())
+				Expect(actualGetResponse).To(BeNil())
+			})
+		})
+
+		When("the get operation can't find the document", func() {
+			BeforeEach(func() {
+				transport.PreparedHttpResponses[0].StatusCode = http.StatusNotFound
+			})
+
+			It("should return the response and no error", func() {
+				Expect(actualErr).ToNot(HaveOccurred())
+				Expect(actualGetResponse).To(Equal(expectedGetResponse))
+			})
+		})
+
+		When("the document id is a url", func() {
+			BeforeEach(func() {
+				expectedDocumentId = fake.URL()
+				expectedGetRequest.DocumentId = expectedDocumentId
+			})
+
+			It("should query escape the document id", func() {
+				Expect(transport.ReceivedHttpRequests[0].Method).To(Equal(http.MethodGet))
+				Expect(transport.ReceivedHttpRequests[0].URL.RawPath).To(ContainSubstring(url.QueryEscape(expectedDocumentId)))
+			})
+		})
+	})
+
 	Context("MultiGet", func() {
 		var (
 			expectedDocumentIds      []string
@@ -625,12 +717,13 @@ var _ = Describe("elasticsearch client", func() {
 			}
 
 			expectedMultiGetResponse = &EsMultiGetResponse{
-				Docs: []*EsMultiGetDocument{},
+				Docs: []*EsGetResponse{},
 			}
 			for _, id := range expectedDocumentIds {
-				expectedMultiGetResponse.Docs = append(expectedMultiGetResponse.Docs, &EsMultiGetDocument{
-					ID:    id,
-					Found: fake.Bool(),
+				expectedMultiGetResponse.Docs = append(expectedMultiGetResponse.Docs, &EsGetResponse{
+					Id:     id,
+					Found:  fake.Bool(),
+					Source: []byte("null"),
 				})
 			}
 
@@ -662,8 +755,10 @@ var _ = Describe("elasticsearch client", func() {
 
 		When("the multiget operation fails", func() {
 			BeforeEach(func() {
-				transport.PreparedHttpResponses[0] = &http.Response{
-					StatusCode: http.StatusInternalServerError,
+				transport.PreparedHttpResponses = []*http.Response{
+					{
+						StatusCode: http.StatusInternalServerError,
+					},
 				}
 			})
 

go/v1beta1/storage/esutil/esutilfakes/fake_client.go

@@ -249,17 +249,18 @@ type ESMeta struct {
 	Type string `json:"type,omitempty"`
 }
 
-type EsMultiGetRequest struct {
-	IDs []string `json:"ids"`
+type EsGetResponse struct {
+	Id     string          `json:"_id"`
+	Found  bool            `json:"found"`
+	Source json.RawMessage `json:"_source"`
 }
 
-type EsMultiGetDocument struct {
-	ID    string `json:"_id"`
-	Found bool   `json:"found"`
+type EsMultiGetRequest struct {
+	IDs []string `json:"ids"`
 }
 
 type EsMultiGetResponse struct {
-	Docs []*EsMultiGetDocument `json:"docs"`
+	Docs []*EsGetResponse `json:"docs"`
 }
 
 // response for index creation