Skip to content

Commit

Permalink
Add json:"" marshaling annotations to response structs (#21)
Browse files Browse the repository at this point in the history
## Problem
There's been some UX feedback regarding marshaling of some of our
response structs.

Primarily, we're not using `json:"key_name"` annotation for our struct
fields so they default to capitalized keys in JSON output when
marshaled. IE: `{ "Usage": { "ReadUnits": 5 }}`. We'd like these lower
case and matched with our API spec.

There are also instances where we're returning pointers for certain
`*int32` fields which can be irritating to consumers needing to
dereference and nil-check. In certain cases we can default any nil
pointer values to 0 before returning things.

There's no testing for the marshaling on any of our custom structs.

## Solution
I may have gone a bit overboard with the testing here. The bulk of this
is tests and the tests are all very similar, definitely open to opinions
on removing if people find them unnecessary, but changing up how the
JSON is produced is pretty easy and it's an important point of
integration in my mind, so maybe worth testing for.

I also wasn't clear if adding these annotations to our request structs
made sense either, so I left those out for now. Overall I'm still
working on getting a better understanding of Go syntax and expectations.

- Add JSON marshaling annotations to the response structs in `models.go`
and `index_connection.go`.
- Update `Collection` and `Usage` structs to default integer values
rather than using pointers.
- Add new `models_test.go` file to run through some validation on struct
marshaling, also add unit tests to `index_connection_test.go`.

Originally I thought `omitempty` for some of these pointers may be
helpful from a consumer standpoint, but I'm actually not sure the more
I've thought about it. Does it make sense to omit more or less things?
Other feedback here would be much appreciated.

## Type of Change
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [X] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [ ] This change requires a documentation update
- [ ] Infrastructure change (CI configs, etc)
- [ ] Non-code change (docs, etc)
- [ ] None of the above: (explain here)

## Test Plan
Unit tests to explicitly stress marshaling the structs via
`json.Marshal()`.
  • Loading branch information
austin-denoble authored May 7, 2024
1 parent 5a9897e commit cbb7f75
Show file tree
Hide file tree
Showing 5 changed files with 837 additions and 67 deletions.
14 changes: 11 additions & 3 deletions pinecone/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -386,12 +386,13 @@ func toCollection(cm *control.CollectionModel) *Collection {
if cm == nil {
return nil
}

return &Collection{
Name: cm.Name,
Size: cm.Size,
Size: derefOrDefault(cm.Size, 0),
Status: CollectionStatus(cm.Status),
Dimension: cm.Dimension,
VectorCount: cm.VectorCount,
Dimension: derefOrDefault(cm.Dimension, 0),
VectorCount: derefOrDefault(cm.VectorCount, 0),
Environment: cm.Environment,
}
}
Expand All @@ -413,6 +414,13 @@ func minOne(x int32) int32 {
return x
}

func derefOrDefault[T any](ptr *T, defaultValue T) T {
if ptr == nil {
return defaultValue
}
return *ptr
}

func buildClientOptions(in NewClientParams) ([]control.ClientOption, error) {
clientOptions := []control.ClientOption{}
hasAuthorizationHeader := false
Expand Down
45 changes: 23 additions & 22 deletions pinecone/index_connection.go
Original file line number Diff line number Diff line change
Expand Up @@ -14,18 +14,18 @@ import (
)

type IndexConnection struct {
Namespace string
apiKey string
Namespace string
apiKey string
additionalMetadata map[string]string
dataClient *data.VectorServiceClient
grpcConn *grpc.ClientConn
dataClient *data.VectorServiceClient
grpcConn *grpc.ClientConn
}

type newIndexParameters struct {
apiKey string
host string
namespace string
sourceTag string
apiKey string
host string
namespace string
sourceTag string
additionalMetadata map[string]string
}

Expand Down Expand Up @@ -75,8 +75,8 @@ func (idx *IndexConnection) UpsertVectors(ctx context.Context, in []*Vector) (ui
}

type FetchVectorsResponse struct {
Vectors map[string]*Vector
Usage *Usage
Vectors map[string]*Vector `json:"vectors,omitempty"`
Usage *Usage `json:"usage,omitempty"`
}

func (idx *IndexConnection) FetchVectors(ctx context.Context, ids []string) (*FetchVectorsResponse, error) {
Expand All @@ -94,6 +94,7 @@ func (idx *IndexConnection) FetchVectors(ctx context.Context, ids []string) (*Fe
for id, vector := range res.Vectors {
vectors[id] = toVector(vector)
}
fmt.Printf("VECTORS: %+v\n", vectors)

return &FetchVectorsResponse{
Vectors: vectors,
Expand All @@ -108,9 +109,9 @@ type ListVectorsRequest struct {
}

type ListVectorsResponse struct {
VectorIds []*string
Usage *Usage
NextPaginationToken *string
VectorIds []*string `json:"vector_ids,omitempty"`
Usage *Usage `json:"usage,omitempty"`
NextPaginationToken *string `json:"next_pagination_token,omitempty"`
}

func (idx *IndexConnection) ListVectors(ctx context.Context, in *ListVectorsRequest) (*ListVectorsResponse, error) {
Expand All @@ -132,7 +133,7 @@ func (idx *IndexConnection) ListVectors(ctx context.Context, in *ListVectorsRequ

return &ListVectorsResponse{
VectorIds: vectorIds,
Usage: &Usage{ReadUnits: res.Usage.ReadUnits},
Usage: &Usage{ReadUnits: derefOrDefault(res.Usage.ReadUnits, 0)},
NextPaginationToken: toPaginationToken(res.Pagination),
}, nil
}
Expand All @@ -147,8 +148,8 @@ type QueryByVectorValuesRequest struct {
}

type QueryVectorsResponse struct {
Matches []*ScoredVector
Usage *Usage
Matches []*ScoredVector `json:"matches,omitempty"`
Usage *Usage `json:"usage,omitempty"`
}

func (idx *IndexConnection) QueryByVectorValues(ctx context.Context, in *QueryByVectorValuesRequest) (*QueryVectorsResponse, error) {
Expand Down Expand Up @@ -236,10 +237,10 @@ func (idx *IndexConnection) UpdateVector(ctx context.Context, in *UpdateVectorRe
}

type DescribeIndexStatsResponse struct {
Dimension uint32
IndexFullness float32
TotalVectorCount uint32
Namespaces map[string]*NamespaceSummary
Dimension uint32 `json:"dimension"`
IndexFullness float32 `json:"index_fullness"`
TotalVectorCount uint32 `json:"total_vector_count"`
Namespaces map[string]*NamespaceSummary `json:"namespaces,omitempty"`
}

func (idx *IndexConnection) DescribeIndexStats(ctx context.Context) (*DescribeIndexStatsResponse, error) {
Expand Down Expand Up @@ -335,7 +336,7 @@ func toUsage(u *data.Usage) *Usage {
return nil
}
return &Usage{
ReadUnits: u.ReadUnits,
ReadUnits: derefOrDefault(u.ReadUnits, 0),
}
}

Expand Down Expand Up @@ -372,7 +373,7 @@ func (idx *IndexConnection) akCtx(ctx context.Context) context.Context {
newMetadata := []string{}
newMetadata = append(newMetadata, "api-key", idx.apiKey)

for key, value := range idx.additionalMetadata{
for key, value := range idx.additionalMetadata {
newMetadata = append(newMetadata, key, value)
}

Expand Down
Loading

0 comments on commit cbb7f75

Please sign in to comment.