Skip to content

Commit

Permalink
Implement IndexTags (#88)
Browse files Browse the repository at this point in the history
## Problem
[Index tags are a feature
](https://docs.pinecone.io/guides/indexes/tag-an-index) which was
included in the `2024-10` API specification, but was not implemented in
the SDK interface.

## Solution
- Add new `IndexTags` type to `models.go`. This just represents
`map[string]string`, lining up with the generated type.
- Update `CreateServerlessIndexRequest`, `CreatePodIndexRequest`, and
`ConfigureIndexParams` to allow passing `IndexTags` optionally.
- Update integration test setup / teardown to support working with tags.
Add a new test method to exercise validating tags, and updating them
through `ConfigureIndex`.

## Type of Change
- [ ] Bug fix (non-breaking change which fixes an issue)
- [X] New feature (non-breaking change which adds functionality)
- [ ] 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
CI Integration Tests should pass - these create serverless & pod indexes
with tags, and then a test checks the tags and updates them as a part of
the integration test suite.

To work with tags directly in Go:

```go
	indexTags := IndexTags{"test1": "test-tag-1", "test2": "test-tag-2"}
	serverlessIdx, err := in.CreateServerlessIndex(ctx, &CreateServerlessIndexRequest{
		Name:      "test-index",
		Dimension: uint32(5),
		Metric:    Cosine,
		Region:    "us-east-1",
		Cloud:     "aws",
		Tags:      &indexTags,
	})
```


---
- To see the specific tasks where the Asana app for GitHub is being
used, see below:
  - https://app.asana.com/0/0/1208820098587324
  • Loading branch information
austin-denoble authored Dec 13, 2024
1 parent 1aa2cd5 commit d406377
Show file tree
Hide file tree
Showing 6 changed files with 74 additions and 11 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/ci.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ jobs:
env:
PORT: 5081
DIMENSION: 1536
METRIC: dot-product
METRIC: dotproduct
INDEX_TYPE: serverless
pc-index-pod:
image: ghcr.io/pinecone-io/pinecone-index:latest
Expand Down
30 changes: 26 additions & 4 deletions pinecone/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -415,6 +415,7 @@ func (c *Client) ListIndexes(ctx context.Context) ([]*Index, error) {
// only valid for use with pod-based Indexes.
// - DeletionProtection: (Optional) determines whether [deletion protection] is "enabled" or "disabled" for the index.
// When "enabled", the index cannot be deleted. Defaults to "disabled".
// - Tags: (Optional) A map of tags to associate with the Index.
//
// To create a new pods-based Index, use the [Client.CreatePodIndex] method.
//
Expand Down Expand Up @@ -473,6 +474,7 @@ type CreatePodIndexRequest struct {
Replicas int32
SourceCollection *string
MetadataConfig *PodSpecMetadataConfig
Tags *IndexTags
}

// [CreatePodIndexRequestReplicaCount] ensures the replica count of a pods-based Index is >1.
Expand Down Expand Up @@ -548,11 +550,17 @@ func (c *Client) CreatePodIndex(ctx context.Context, in *CreatePodIndexRequest)
replicas := in.ReplicaCount()
shards := in.ShardCount()

var tags *db_control.IndexTags
if in.Tags != nil {
tags = (*db_control.IndexTags)(in.Tags)
}

req := db_control.CreateIndexRequest{
Name: in.Name,
Dimension: in.Dimension,
Metric: metric,
DeletionProtection: deletionProtection,
Tags: tags,
}

req.Spec = db_control.IndexSpec{
Expand Down Expand Up @@ -595,11 +603,12 @@ func (c *Client) CreatePodIndex(ctx context.Context, in *CreatePodIndexRequest)
// and consist only of lower case alphanumeric characters or '-'.
// - Dimension: (Required) The [dimensionality] of the vectors to be inserted in the [Index].
// - Metric: (Required) The metric used to measure the [similarity] between vectors ('euclidean', 'cosine', or 'dotproduct').
// - DeletionProtection: (Optional) Determines whether [deletion protection] is "enabled" or "disabled" for the index.
// When "enabled", the index cannot be deleted. Defaults to "disabled".
// - Cloud: (Required) The public [cloud provider] where you would like your [Index] hosted.
// For serverless Indexes, you define only the cloud and region where the [Index] should be hosted.
// - Region: (Required) The [region] where you would like your [Index] to be created.
// - DeletionProtection: (Optional) Determines whether [deletion protection] is "enabled" or "disabled" for the index.
// When "enabled", the index cannot be deleted. Defaults to "disabled".
// - Tags: (Optional) A map of tags to associate with the Index.
//
// To create a new Serverless Index, use the [Client.CreateServerlessIndex] method.
//
Expand Down Expand Up @@ -648,6 +657,7 @@ type CreateServerlessIndexRequest struct {
DeletionProtection DeletionProtection
Cloud Cloud
Region string
Tags *IndexTags
}

// [Client.CreateServerlessIndex] creates and initializes a new serverless Index via the specified [Client].
Expand Down Expand Up @@ -698,6 +708,11 @@ func (c *Client) CreateServerlessIndex(ctx context.Context, in *CreateServerless
deletionProtection := pointerOrNil(db_control.DeletionProtection(in.DeletionProtection))
metric := pointerOrNil(db_control.CreateIndexRequestMetric(in.Metric))

var tags *db_control.IndexTags
if in.Tags != nil {
tags = (*db_control.IndexTags)(in.Tags)
}

req := db_control.CreateIndexRequest{
Name: in.Name,
Dimension: in.Dimension,
Expand All @@ -709,6 +724,7 @@ func (c *Client) CreateServerlessIndex(ctx context.Context, in *CreateServerless
Region: in.Region,
},
},
Tags: tags,
}

res, err := c.restClient.CreateIndex(ctx, req)
Expand Down Expand Up @@ -838,6 +854,7 @@ func (c *Client) DeleteIndex(ctx context.Context, idxName string) error {
// go to [app.pinecone.io], select your project, and configure the maximum number of pods.
// - DeletionProtection: (Optional) DeletionProtection determines whether [deletion protection]
// is "enabled" or "disabled" for the index. When "enabled", the index cannot be deleted. Defaults to "disabled".
// - Tags: (Optional) A map of tags to associate with the Index.
//
// Example:
//
Expand All @@ -864,6 +881,7 @@ type ConfigureIndexParams struct {
PodType string
Replicas int32
DeletionProtection DeletionProtection
Tags IndexTags
}

// [Client.ConfigureIndex] is used to [scale a pods-based index] up or down by changing the size of the pods or the number of
Expand Down Expand Up @@ -908,8 +926,8 @@ type ConfigureIndexParams struct {
//
// [scale a pods-based index]: https://docs.pinecone.io/guides/indexes/configure-pod-based-indexes
func (c *Client) ConfigureIndex(ctx context.Context, name string, in ConfigureIndexParams) (*Index, error) {
if in.PodType == "" && in.Replicas == 0 && in.DeletionProtection == "" {
return nil, fmt.Errorf("must specify PodType, Replicas, or DeletionProtection when configuring an index")
if in.PodType == "" && in.Replicas == 0 && in.DeletionProtection == "" && in.Tags == nil {
return nil, fmt.Errorf("must specify PodType, Replicas, DeletionProtection, or Tags when configuring an index")
}

podType := pointerOrNil(in.PodType)
Expand All @@ -935,6 +953,8 @@ func (c *Client) ConfigureIndex(ctx context.Context, name string, in ConfigureIn
}
}
request.DeletionProtection = (*db_control.DeletionProtection)(deletionProtection)
request.Tags = (*db_control.IndexTags)(&in.Tags)
fmt.Printf("request.Tags: %+v\n", request.Tags)

res, err := c.restClient.ConfigureIndex(ctx, name, request)
if err != nil {
Expand Down Expand Up @@ -1528,6 +1548,7 @@ func toIndex(idx *db_control.IndexModel) *Index {
Ready: idx.Status.Ready,
State: IndexStatusState(idx.Status.State),
}
tags := (*IndexTags)(idx.Tags)
deletionProtection := derefOrDefault(idx.DeletionProtection, "disabled")

return &Index{
Expand All @@ -1538,6 +1559,7 @@ func toIndex(idx *db_control.IndexModel) *Index {
DeletionProtection: DeletionProtection(deletionProtection),
Spec: spec,
Status: status,
Tags: tags,
}
}

Expand Down
35 changes: 33 additions & 2 deletions pinecone/client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -131,10 +131,10 @@ func (ts *IntegrationTests) TestListCollections() {
}

func (ts *IntegrationTests) TestDescribeCollection() {
ctx := context.Background()
if ts.indexType == "serverless" {
ts.T().Skip("No pod index to test")
}
ctx := context.Background()

collection, err := ts.client.DescribeCollection(ctx, ts.collectionName)
require.NoError(ts.T(), err)
Expand Down Expand Up @@ -250,7 +250,7 @@ func (ts *IntegrationTests) TestConfigureIndexScaleUpNoReplicas() {

func (ts *IntegrationTests) TestConfigureIndexIllegalNoPodsOrReplicasOrDeletionProtection() {
_, err := ts.client.ConfigureIndex(context.Background(), ts.idxName, ConfigureIndexParams{})
require.ErrorContainsf(ts.T(), err, "must specify PodType, Replicas, or DeletionProtection", err.Error())
require.ErrorContainsf(ts.T(), err, "must specify PodType, Replicas, DeletionProtection, or Tags", err.Error())
}

func (ts *IntegrationTests) TestConfigureIndexHitPodLimit() {
Expand Down Expand Up @@ -476,6 +476,37 @@ func (ts *IntegrationTests) TestRerankDocumentFieldError() {
require.Contains(ts.T(), err.Error(), "field 'custom-field' not found in document")
}

func (ts *IntegrationTests) TestIndexTags() {
// Validate that index tags are set
index, err := ts.client.DescribeIndex(context.Background(), ts.idxName)
require.NoError(ts.T(), err)

assert.Equal(ts.T(), ts.indexTags, index.Tags, "Expected index tags to match")

// Update first tag, and clear the second
counter := 0
updatedTags := make(IndexTags)
deletedTag := ""
for key := range *ts.indexTags {
if counter == 0 {
updatedTags[key] = "updated-tag"
} else {
deletedTag = key
updatedTags[key] = ""
}
counter++
}

index, err = ts.client.ConfigureIndex(context.Background(), ts.idxName, ConfigureIndexParams{Tags: updatedTags})
require.NoError(ts.T(), err)

// Remove empty tag from the map
delete(updatedTags, deletedTag)

assert.Equal(ts.T(), &updatedTags, index.Tags, "Expected index tags to match")
ts.indexTags = &updatedTags
}

// Unit tests:
func TestExtractAuthHeaderUnit(t *testing.T) {
globalApiKey := os.Getenv("PINECONE_API_KEY")
Expand Down
4 changes: 4 additions & 0 deletions pinecone/models.go
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,9 @@ type IndexSpec struct {
Serverless *ServerlessSpec `json:"serverless,omitempty"`
}

// [IndexTags] is a set of key-value pairs that can be attached to a Pinecone [Index].
type IndexTags map[string]string

// [Index] is a Pinecone [Index] object. Can be either a pod-based or a serverless [Index], depending on the [IndexSpec].
type Index struct {
Name string `json:"name"`
Expand All @@ -74,6 +77,7 @@ type Index struct {
DeletionProtection DeletionProtection `json:"deletion_protection,omitempty"`
Spec *IndexSpec `json:"spec,omitempty"`
Status *IndexStatus `json:"status,omitempty"`
Tags *IndexTags `json:"tags,omitempty"`
}

// [Collection] is a Pinecone [collection entity]. Only available for pod-based Indexes.
Expand Down
7 changes: 5 additions & 2 deletions pinecone/suite_runner_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,10 @@ func RunSuites(t *testing.T) {
client, err := NewClient(NewClientParams{ApiKey: apiKey, SourceTag: sourceTag})
require.NotNil(t, client, "Client should not be nil after creation")
require.NoError(t, err)
indexTags := IndexTags{"test1": "test-tag-1", "test2": "test-tag-2"}

serverlessIdx := BuildServerlessTestIndex(client, "serverless-"+GenerateTestIndexName())
podIdx := BuildPodTestIndex(client, "pods-"+GenerateTestIndexName())
serverlessIdx := BuildServerlessTestIndex(client, "serverless-"+GenerateTestIndexName(), indexTags)
podIdx := BuildPodTestIndex(client, "pods-"+GenerateTestIndexName(), indexTags)

podTestSuite := &IntegrationTests{
apiKey: apiKey,
Expand All @@ -36,6 +37,7 @@ func RunSuites(t *testing.T) {
client: client,
sourceTag: sourceTag,
idxName: podIdx.Name,
indexTags: &indexTags,
}

serverlessTestSuite := &IntegrationTests{
Expand All @@ -46,6 +48,7 @@ func RunSuites(t *testing.T) {
client: client,
sourceTag: sourceTag,
idxName: serverlessIdx.Name,
indexTags: &indexTags,
}

suite.Run(t, podTestSuite)
Expand Down
7 changes: 5 additions & 2 deletions pinecone/test_suite.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ type IntegrationTests struct {
idxConn *IndexConnection
collectionName string
sourceTag string
indexTags *IndexTags
}

func (ts *IntegrationTests) SetupSuite() {
Expand Down Expand Up @@ -207,7 +208,7 @@ func generateVectorValues(dimension int32) []float32 {
return values
}

func BuildServerlessTestIndex(in *Client, idxName string) *Index {
func BuildServerlessTestIndex(in *Client, idxName string, tags IndexTags) *Index {
ctx := context.Background()

fmt.Printf("Creating Serverless index: %s\n", idxName)
Expand All @@ -217,6 +218,7 @@ func BuildServerlessTestIndex(in *Client, idxName string) *Index {
Metric: Cosine,
Region: "us-east-1",
Cloud: "aws",
Tags: &tags,
})
if err != nil {
log.Fatalf("Failed to create Serverless index \"%s\" in integration test: %v", err, idxName)
Expand All @@ -226,7 +228,7 @@ func BuildServerlessTestIndex(in *Client, idxName string) *Index {
return serverlessIdx
}

func BuildPodTestIndex(in *Client, name string) *Index {
func BuildPodTestIndex(in *Client, name string, tags IndexTags) *Index {
ctx := context.Background()

fmt.Printf("Creating pod index: %s\n", name)
Expand All @@ -236,6 +238,7 @@ func BuildPodTestIndex(in *Client, name string) *Index {
Metric: Cosine,
Environment: "us-east-1-aws",
PodType: "p1",
Tags: &tags,
})
if err != nil {
log.Fatalf("Failed to create pod index in buildPodTestIndex test: %v", err)
Expand Down

0 comments on commit d406377

Please sign in to comment.