Skip to content

Commit

Permalink
review feedback, fix doc comments, add two integration tests for dele…
Browse files Browse the repository at this point in the history 
…tion protection
  • Loading branch information
@austin-denoble
austin-denoble committed Jul 26, 2024
1 parent f66ca9d commit 6677b88
Show file tree
Hide file tree
Showing 3 changed files with 108 additions and 32 deletions.
59 changes: 35 additions & 24 deletions pinecone/client.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -383,6 +383,8 @@ func (c *Client) ListIndexes(ctx context.Context) ([]*Index, error) {
// default, all metadata is indexed; when `metadata_config` is present,
// only specified metadata fields are indexed. These configurations are
// 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".
//
// To create a new pods-based Index, use the CreatePodIndex method on the Client object.
//
Expand Down Expand Up @@ -426,8 +428,8 @@ func (c *Client) ListIndexes(ctx context.Context) ([]*Index, error) {
// [metadata configuration]: https://docs.pinecone.io/guides/indexes/configure-pod-based-indexes#selective-metadata-indexing
// [cloud environment]: https://docs.pinecone.io/guides/indexes/understanding-indexes#pod-environments
// [replicas]: https://docs.pinecone.io/guides/indexes/configure-pod-based-indexes#add-replicas
//
// [type of pods]: https://docs.pinecone.io/guides/indexes/choose-a-pod-type-and-size
// [type of pod]: https://docs.pinecone.io/guides/indexes/choose-a-pod-type-and-size
// [deletion protection]: https://docs.pinecone.io/guides/indexes/prevent-index-deletion#enable-deletion-protection
type CreatePodIndexRequest struct {
Name string
Dimension int32
Expand Down Expand Up @@ -552,10 +554,11 @@ func (c *Client) CreatePodIndex(ctx context.Context, in *CreatePodIndexRequest)
// and consist only of lower case alphanumeric characters or '-'.
// - Dimension: The [dimensionality] of the vectors to be inserted in the Index.
// - Metric: The metric used to measure the [similarity] between vectors ('euclidean', 'cosine', or 'dotproduct').
// - DeletionProtection:
// - Cloud: 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: 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".
//
// To create a new Serverless Index, use the CreateServerlessIndex method on the Client object.
//
Expand Down Expand Up @@ -594,6 +597,7 @@ func (c *Client) CreatePodIndex(ctx context.Context, in *CreatePodIndexRequest)
// [similarity]: https://docs.pinecone.io/guides/indexes/understanding-indexes#distance-metrics
// [region]: https://docs.pinecone.io/troubleshooting/available-cloud-regions
// [cloud provider]: https://docs.pinecone.io/troubleshooting/available-cloud-regions#regions-available-for-serverless-indexes
// [deletion protection]: https://docs.pinecone.io/guides/indexes/prevent-index-deletion#enable-deletion-protection
type CreateServerlessIndexRequest struct {
Name string
Dimension int32
Expand Down Expand Up @@ -776,7 +780,7 @@ func (c *Client) DeleteIndex(ctx context.Context, idxName string) error {
// - Replicas: (Optional) The number of replicas to scale the index to. This is capped by
// the maximum number of replicas allowed in your Pinecone project. To configure this number,
// go to [app.pinecone.io], select your project, and configure the maximum number of pods.
// - DeletionProtection: (Optional) // DeletionProtection determines whether [deletion protection]
// - DeletionProtection: (Optional) DeletionProtection determines whether [deletion protection]
// is "enabled" or "disabled" for the index. When "enabled", the index cannot be deleted. Defaults to "disabled".
//
// Example:
Expand All @@ -799,10 +803,11 @@ func (c *Client) DeleteIndex(ctx context.Context, idxName string) error {
//
// [app.pinecone.io]: https://app.pinecone.io
// [scale a pods-based index]: https://docs.pinecone.io/guides/indexes/configure-pod-based-indexes
// [deletion protection]: https://docs.pinecone.io/guides/indexes/prevent-index-deletion#enable-deletion-protection
type ConfigureIndexParams struct {
PodType string
Replicas int32
DeletionProtection control.DeletionProtection
DeletionProtection DeletionProtection
}

// 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 @@ -846,33 +851,34 @@ 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) {

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 either PodType, Replicas, or DeletionProtection")
return nil, fmt.Errorf("must specify PodType, Replicas, or DeletionProtection when configuring an index")
}

podType := pointerOrNil(in.PodType)
replicas := pointerOrNil(in.Replicas)
deletionProtection := pointerOrNil(in.DeletionProtection)

request := control.ConfigureIndexRequest{
Spec: &struct {
Pod struct {
PodType *string `json:"pod_type,omitempty"`
Replicas *int32 `json:"replicas,omitempty"`
} `json:"pod"`
}{
Pod: struct {
PodType *string `json:"pod_type,omitempty"`
Replicas *int32 `json:"replicas,omitempty"`
var request control.ConfigureIndexRequest
if podType != nil || replicas != nil {
request.Spec =
&struct {
Pod struct {
PodType *string `json:"pod_type,omitempty"`
Replicas *int32 `json:"replicas,omitempty"`
} `json:"pod"`
}{
PodType: podType,
Replicas: replicas,
},
},
DeletionProtection: deletionProtection,
Pod: struct {
PodType *string `json:"pod_type,omitempty"`
Replicas *int32 `json:"replicas,omitempty"`
}{
PodType: podType,
Replicas: replicas,
},
}
}
request.DeletionProtection = (*control.DeletionProtection)(deletionProtection)

res, err := c.restClient.ConfigureIndex(ctx, name, request)
if err != nil {
Expand All @@ -888,6 +894,11 @@ func (c *Client) ConfigureIndex(ctx context.Context, name string, in *ConfigureI
return decodeIndex(res.Body)
}

func PrettifyStruct(obj interface{}) string {
bytes, _ := json.MarshalIndent(obj, "", " ")
return string(bytes)
}

// ListCollections retrieves a list of all Collections in a Pinecone [project]. See Collection for more information.
//
// Parameters:
Expand Down Expand Up @@ -1209,7 +1220,7 @@ func toIndex(idx *control.IndexModel) *Index {
Ready: idx.Status.Ready,
State: IndexStatusState(idx.Status.State),
}
deletionProtection := derefOrDefault(idx.DeletionProtection, "")
deletionProtection := derefOrDefault(idx.DeletionProtection, "disabled")

return &Index{
Name: idx.Name,
Expand Down
77 changes: 71 additions & 6 deletions pinecone/client_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -465,6 +465,71 @@ func (ts *ClientTestsIntegration) TestDeleteCollection() {
require.NoError(ts.T(), err)
}

func (ts *ClientTestsIntegration) TestDeletionProtectionPodIntegration() {
name := uuid.New().String()

defer func(ts *ClientTestsIntegration, name string) {
err := ts.deleteIndex(name)
require.NoError(ts.T(), err)
}(ts, name)

_, err := ts.client.CreatePodIndex(context.Background(), &CreatePodIndexRequest{
Name: name,
Dimension: 10,
Metric: Cosine,
Environment: "us-east1-gcp",
PodType: "p1.x1",
DeletionProtection: "enabled",
Shards: 1,
Replicas: 1,
})
require.NoError(ts.T(), err)

// validate we cannot delete the index
err = ts.client.DeleteIndex(context.Background(), name)
require.ErrorContainsf(ts.T(), err, "failed to delete index: Deletion protection is enabled for this index", err.Error())

// make sure configuring the index without specifying DeletionProtection maintains "enabled" state
index, err := ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{PodType: "p1.x2"})
require.NoError(ts.T(), err)
require.Equal(ts.T(), DeletionProtectionEnabled, index.DeletionProtection, "Expected deletion protection to be 'enabled'")

// validate we cannot delete the index
err = ts.client.DeleteIndex(context.Background(), name)
require.ErrorContainsf(ts.T(), err, "failed to delete index: Deletion protection is enabled for this index", err.Error())

// disable deletion protection so we can clean up
_, err = ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{DeletionProtection: "disabled"})
require.NoError(ts.T(), err)
}

func (ts *ClientTestsIntegration) TestDeletionProtectionServerlessIntegration() {
name := uuid.New().String()

defer func(ts *ClientTestsIntegration, name string) {
err := ts.deleteIndex(name)
require.NoError(ts.T(), err)
}(ts, name)

_, err := ts.client.CreateServerlessIndex(context.Background(), &CreateServerlessIndexRequest{
Name: name,
Dimension: 10,
Metric: Cosine,
Cloud: "aws",
Region: "us-west-2",
DeletionProtection: "enabled",
})
require.NoError(ts.T(), err)

// validate we cannot delete the index
err = ts.client.DeleteIndex(context.Background(), name)
require.ErrorContainsf(ts.T(), err, "failed to delete index: Deletion protection is enabled for this index", err.Error())

// disable deletion protection so we can clean up
_, err = ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{DeletionProtection: "disabled"})
require.NoError(ts.T(), err)
}

func (ts *ClientTestsIntegration) TestConfigureIndexIllegalScaleDown() {
name := uuid.New().String()

Expand All @@ -484,7 +549,7 @@ func (ts *ClientTestsIntegration) TestConfigureIndexIllegalScaleDown() {
log.Fatalf("Error creating index %s: %v", name, err)
}

_, err = ts.client.ConfigureIndex(context.Background(), name, &ConfigureIndexParams{PodType: "p1.x1", Replicas: 1})
_, err = ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{PodType: "p1.x1"})
require.ErrorContainsf(ts.T(), err, "Cannot scale down", err.Error())
}

Expand All @@ -507,7 +572,7 @@ func (ts *ClientTestsIntegration) TestConfigureIndexScaleUpNoPods() {
log.Fatalf("Error creating index %s: %v", name, err)
}

_, err = ts.client.ConfigureIndex(context.Background(), name, &ConfigureIndexParams{Replicas: 2})
_, err = ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{Replicas: 2})
require.NoError(ts.T(), err)
}

Expand All @@ -530,7 +595,7 @@ func (ts *ClientTestsIntegration) TestConfigureIndexScaleUpNoReplicas() {
log.Fatalf("Error creating index %s: %v", name, err)
}

_, err = ts.client.ConfigureIndex(context.Background(), name, &ConfigureIndexParams{PodType: "p1.x4"})
_, err = ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{PodType: "p1.x4"})
require.NoError(ts.T(), err)
}

Expand All @@ -553,8 +618,8 @@ func (ts *ClientTestsIntegration) TestConfigureIndexIllegalNoPodsOrReplicasOrDel
log.Fatalf("Error creating index %s: %v", name, err)
}

_, err = ts.client.ConfigureIndex(context.Background(), name, &ConfigureIndexParams{})
require.ErrorContainsf(ts.T(), err, "must specify either PodType, Replicas, or DeletionProtection", err.Error())
_, err = ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{})
require.ErrorContainsf(ts.T(), err, "must specify PodType, Replicas, or DeletionProtection", err.Error())
}

func (ts *ClientTestsIntegration) TestConfigureIndexHitPodLimit() {
Expand All @@ -576,7 +641,7 @@ func (ts *ClientTestsIntegration) TestConfigureIndexHitPodLimit() {
log.Fatalf("Error creating index %s: %v", name, err)
}

_, err = ts.client.ConfigureIndex(context.Background(), name, &ConfigureIndexParams{Replicas: 30000})
_, err = ts.client.ConfigureIndex(context.Background(), name, ConfigureIndexParams{Replicas: 30000})
require.ErrorContainsf(ts.T(), err, "You've reached the max pods allowed", err.Error())
}

Expand Down
4 changes: 2 additions & 2 deletions pinecone/models.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,8 +36,8 @@ const (
type DeletionProtection string

const (
DeletionProtectionEnabled DeletionProtection = "Enabled"
DeletionProtectionDisabled DeletionProtection = "Disabled"
DeletionProtectionEnabled DeletionProtection = "enabled"
DeletionProtectionDisabled DeletionProtection = "disabled"
)

// Cloud is the [cloud provider] to be used for a Pinecone serverless Index.
Expand Down

0 comments on commit 6677b88

Please sign in to comment.