diff --git a/internal/pkg/cli/command/apiKey/create.go b/internal/pkg/cli/command/apiKey/create.go index f7e84fe..d614870 100644 --- a/internal/pkg/cli/command/apiKey/create.go +++ b/internal/pkg/cli/command/apiKey/create.go @@ -57,7 +57,8 @@ func NewCreateApiKeyCmd() *cobra.Command { pc api-key create --id "project-id" --name "key-name" `), Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) projId := options.projectId var err error diff --git a/internal/pkg/cli/command/apiKey/delete.go b/internal/pkg/cli/command/apiKey/delete.go index 65543db..a5b0dda 100644 --- a/internal/pkg/cli/command/apiKey/delete.go +++ b/internal/pkg/cli/command/apiKey/delete.go @@ -46,7 +46,8 @@ func NewDeleteKeyCmd() *cobra.Command { pc api-key delete --id "api-key-id" `), Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) // Verify key exists before trying to delete it. // This lets us give a more helpful error message than just diff --git a/internal/pkg/cli/command/apiKey/describe.go b/internal/pkg/cli/command/apiKey/describe.go index 2be9e1c..55922c4 100644 --- a/internal/pkg/cli/command/apiKey/describe.go +++ b/internal/pkg/cli/command/apiKey/describe.go @@ -28,7 +28,8 @@ func NewDescribeAPIKeyCmd() *cobra.Command { `), GroupID: help.GROUP_API_KEYS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) apiKey, err := ac.APIKey.Describe(cmd.Context(), options.apiKeyID) if err != nil { diff --git a/internal/pkg/cli/command/apiKey/list.go b/internal/pkg/cli/command/apiKey/list.go index 6d42e13..1d41c5b 100644 --- a/internal/pkg/cli/command/apiKey/list.go +++ b/internal/pkg/cli/command/apiKey/list.go @@ -38,7 +38,8 @@ func NewListKeysCmd() *cobra.Command { `), GroupID: help.GROUP_API_KEYS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) var err error projId := options.projectID diff --git a/internal/pkg/cli/command/apiKey/update.go b/internal/pkg/cli/command/apiKey/update.go index cd469c8..5338737 100644 --- a/internal/pkg/cli/command/apiKey/update.go +++ b/internal/pkg/cli/command/apiKey/update.go @@ -30,7 +30,8 @@ func NewUpdateAPIKeyCmd() *cobra.Command { `), GroupID: help.GROUP_API_KEYS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) // Only set non-empty values updateParams := &pinecone.UpdateAPIKeyParams{} diff --git a/internal/pkg/cli/command/auth/configure.go b/internal/pkg/cli/command/auth/configure.go index 6a9bdd0..7b9b72a 100644 --- a/internal/pkg/cli/command/auth/configure.go +++ b/internal/pkg/cli/command/auth/configure.go @@ -139,7 +139,7 @@ func Run(ctx context.Context, io IO, opts configureCmdOptions) { // Use Admin API to fetch organization and project information for the service account // so that we can set the target context, or allow the user to set it like they do through the login or target flow - ac := sdk.NewPineconeAdminClient() + ac := sdk.NewPineconeAdminClient(ctx) // There should only be one organization listed for a service account orgs, err := ac.Organization.List(ctx) diff --git a/internal/pkg/cli/command/auth/local_keys_prune.go b/internal/pkg/cli/command/auth/local_keys_prune.go index 6098878..26e489b 100644 --- a/internal/pkg/cli/command/auth/local_keys_prune.go +++ b/internal/pkg/cli/command/auth/local_keys_prune.go @@ -80,7 +80,7 @@ func NewPruneLocalKeysCmd() *cobra.Command { } func runPruneLocalKeys(ctx context.Context, options pruneLocalKeysCmdOptions) { - ac := sdk.NewPineconeAdminClient() + ac := sdk.NewPineconeAdminClient(ctx) managedKeys := secrets.GetManagedProjectKeys() // Filter to projectId if provided diff --git a/internal/pkg/cli/command/collection/create.go b/internal/pkg/cli/command/collection/create.go index f145ab3..16307f5 100644 --- a/internal/pkg/cli/command/collection/create.go +++ b/internal/pkg/cli/command/collection/create.go @@ -1,8 +1,6 @@ package collection import ( - "context" - "github.com/pinecone-io/cli/internal/pkg/utils/exit" "github.com/pinecone-io/cli/internal/pkg/utils/help" "github.com/pinecone-io/cli/internal/pkg/utils/msg" @@ -32,8 +30,8 @@ func NewCreateCollectionCmd() *cobra.Command { pc collection create --name "collection-name" --source "index-source-name" `), Run: func(cmd *cobra.Command, args []string) { - pc := sdk.NewPineconeClient() - ctx := context.Background() + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) req := &pinecone.CreateCollectionRequest{ Name: options.name, diff --git a/internal/pkg/cli/command/collection/delete.go b/internal/pkg/cli/command/collection/delete.go index 1f78eaf..9dea8da 100644 --- a/internal/pkg/cli/command/collection/delete.go +++ b/internal/pkg/cli/command/collection/delete.go @@ -1,8 +1,6 @@ package collection import ( - "context" - "github.com/pinecone-io/cli/internal/pkg/utils/exit" "github.com/pinecone-io/cli/internal/pkg/utils/help" "github.com/pinecone-io/cli/internal/pkg/utils/msg" @@ -25,8 +23,8 @@ func NewDeleteCollectionCmd() *cobra.Command { pc collection delete --name "collection-name" `), Run: func(cmd *cobra.Command, args []string) { - ctx := context.Background() - pc := sdk.NewPineconeClient() + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) err := pc.DeleteCollection(ctx, options.name) if err != nil { diff --git a/internal/pkg/cli/command/collection/describe.go b/internal/pkg/cli/command/collection/describe.go index dd6ae19..ac774bb 100644 --- a/internal/pkg/cli/command/collection/describe.go +++ b/internal/pkg/cli/command/collection/describe.go @@ -1,8 +1,6 @@ package collection import ( - "context" - "github.com/pinecone-io/cli/internal/pkg/utils/exit" "github.com/pinecone-io/cli/internal/pkg/utils/help" "github.com/pinecone-io/cli/internal/pkg/utils/msg" @@ -28,8 +26,8 @@ func NewDescribeCollectionCmd() *cobra.Command { pc collection describe --name "collection-name" `), Run: func(cmd *cobra.Command, args []string) { - ctx := context.Background() - pc := sdk.NewPineconeClient() + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) collection, err := pc.DescribeCollection(ctx, options.name) if err != nil { diff --git a/internal/pkg/cli/command/collection/list.go b/internal/pkg/cli/command/collection/list.go index a2d41bc..9e28e8f 100644 --- a/internal/pkg/cli/command/collection/list.go +++ b/internal/pkg/cli/command/collection/list.go @@ -1,7 +1,6 @@ package collection import ( - "context" "os" "sort" "strconv" @@ -33,8 +32,8 @@ func NewListCollectionsCmd() *cobra.Command { pc collection list `), Run: func(cmd *cobra.Command, args []string) { - pc := sdk.NewPineconeClient() - ctx := context.Background() + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) collections, err := pc.ListCollections(ctx) if err != nil { diff --git a/internal/pkg/cli/command/index/cmd.go b/internal/pkg/cli/command/index/cmd.go index 284eb3c..a0a8293 100644 --- a/internal/pkg/cli/command/index/cmd.go +++ b/internal/pkg/cli/command/index/cmd.go @@ -39,6 +39,13 @@ func NewIndexCmd() *cobra.Command { cmd.AddCommand(NewCreatePodCmd()) cmd.AddCommand(NewConfigureIndexCmd()) cmd.AddCommand(NewDeleteCmd()) + cmd.AddCommand(NewUpsertCmd()) + cmd.AddCommand(NewFetchCmd()) + cmd.AddCommand(NewQueryCmd()) + cmd.AddCommand(NewListVectorsCmd()) + cmd.AddCommand(NewDeleteVectorsCmd()) + cmd.AddCommand(NewUpdateCmd()) + cmd.AddCommand(NewDescribeIndexStatsCmd()) return cmd } diff --git a/internal/pkg/cli/command/index/configure.go b/internal/pkg/cli/command/index/configure.go index 85358e6..6c10c7e 100644 --- a/internal/pkg/cli/command/index/configure.go +++ b/internal/pkg/cli/command/index/configure.go @@ -33,7 +33,7 @@ func NewConfigureIndexCmd() *cobra.Command { pc index configure --name "index-name" --deletion-protection "enabled" `), Run: func(cmd *cobra.Command, args []string) { - runConfigureIndexCmd(options) + runConfigureIndexCmd(cmd.Context(), options) }, } @@ -48,9 +48,8 @@ func NewConfigureIndexCmd() *cobra.Command { return cmd } -func runConfigureIndexCmd(options configureIndexOptions) { - ctx := context.Background() - pc := sdk.NewPineconeClient() +func runConfigureIndexCmd(ctx context.Context, options configureIndexOptions) { + pc := sdk.NewPineconeClient(ctx) idx, err := pc.ConfigureIndex(ctx, options.name, pinecone.ConfigureIndexParams{ PodType: options.podType, diff --git a/internal/pkg/cli/command/index/create.go b/internal/pkg/cli/command/index/create.go index c01e083..e2f0054 100644 --- a/internal/pkg/cli/command/index/create.go +++ b/internal/pkg/cli/command/index/create.go @@ -99,7 +99,7 @@ func NewCreateIndexCmd() *cobra.Command { Long: createIndexHelp, Example: createIndexExample, Run: func(cmd *cobra.Command, args []string) { - runCreateIndexCmd(options) + runCreateIndexCmd(cmd.Context(), options) }, } @@ -141,9 +141,8 @@ func NewCreateIndexCmd() *cobra.Command { return cmd } -func runCreateIndexCmd(options createIndexOptions) { - ctx := context.Background() - pc := sdk.NewPineconeClient() +func runCreateIndexCmd(ctx context.Context, options createIndexOptions) { + pc := sdk.NewPineconeClient(ctx) idx, err := runCreateIndexWithService(ctx, pc, options) if err != nil { diff --git a/internal/pkg/cli/command/index/create_pod.go b/internal/pkg/cli/command/index/create_pod.go index dfa19a3..da38cf6 100644 --- a/internal/pkg/cli/command/index/create_pod.go +++ b/internal/pkg/cli/command/index/create_pod.go @@ -39,7 +39,7 @@ func NewCreatePodCmd() *cobra.Command { pc index create-pod --name "my-index" --dimension 1536 --metric "cosine" --environment "us-east-1-aws" --pod-type "p1.x1" --shards 2 --replicas 2 `), Run: func(cmd *cobra.Command, args []string) { - runCreatePodCmd(options) + runCreatePodCmd(cmd.Context(), options) }, } @@ -65,9 +65,8 @@ func NewCreatePodCmd() *cobra.Command { return cmd } -func runCreatePodCmd(options createPodOptions) { - ctx := context.Background() - pc := sdk.NewPineconeClient() +func runCreatePodCmd(ctx context.Context, options createPodOptions) { + pc := sdk.NewPineconeClient(ctx) // Deprecation warning pcio.Fprintf(os.Stderr, "⚠️ Warning: The '%s' command is deprecated. Please use '%s' instead.", style.Code("index create-pod"), style.Code("index create")) diff --git a/internal/pkg/cli/command/index/create_serverless.go b/internal/pkg/cli/command/index/create_serverless.go index cfa970f..eb9c9bb 100644 --- a/internal/pkg/cli/command/index/create_serverless.go +++ b/internal/pkg/cli/command/index/create_serverless.go @@ -36,7 +36,7 @@ func NewCreateServerlessCmd() *cobra.Command { pc index create-serverless --name "my-index" --dimension 1536 --metric "cosine" --cloud "aws" --region "us-east-1" `), Run: func(cmd *cobra.Command, args []string) { - runCreateServerlessCmd(options) + runCreateServerlessCmd(cmd.Context(), options) }, } @@ -58,9 +58,8 @@ func NewCreateServerlessCmd() *cobra.Command { return cmd } -func runCreateServerlessCmd(options createServerlessOptions) { - ctx := context.Background() - pc := sdk.NewPineconeClient() +func runCreateServerlessCmd(ctx context.Context, options createServerlessOptions) { + pc := sdk.NewPineconeClient(ctx) // Deprecation warning pcio.Fprintf(os.Stderr, "⚠️ Warning: The '%s' command is deprecated. Please use '%s' instead.", style.Code("index create-serverless"), style.Code("index create")) diff --git a/internal/pkg/cli/command/index/delete.go b/internal/pkg/cli/command/index/delete.go index ae57c00..902da3b 100644 --- a/internal/pkg/cli/command/index/delete.go +++ b/internal/pkg/cli/command/index/delete.go @@ -1,7 +1,6 @@ package index import ( - "context" "strings" "github.com/pinecone-io/cli/internal/pkg/utils/exit" @@ -26,8 +25,8 @@ func NewDeleteCmd() *cobra.Command { pc index delete --name "index-name" `), Run: func(cmd *cobra.Command, args []string) { - ctx := context.Background() - pc := sdk.NewPineconeClient() + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) err := pc.DeleteIndex(ctx, options.name) if err != nil { diff --git a/internal/pkg/cli/command/index/delete_vectors.go b/internal/pkg/cli/command/index/delete_vectors.go new file mode 100644 index 0000000..0c96e28 --- /dev/null +++ b/internal/pkg/cli/command/index/delete_vectors.go @@ -0,0 +1,109 @@ +package index + +import ( + "context" + + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/flags" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/go-pinecone/v5/pinecone" + "github.com/spf13/cobra" +) + +type deleteVectorsCmdOptions struct { + indexName string + namespace string + ids flags.StringList + filter flags.JSONObject + deleteAllVectors bool + json bool +} + +func NewDeleteVectorsCmd() *cobra.Command { + options := deleteVectorsCmdOptions{} + + cmd := &cobra.Command{ + Use: "delete-vectors", + Short: "Delete vectors from an index", + Example: help.Examples(` + pc index delete-vectors --index-name my-index --namespace my-namespace --ids my-id + pc index delete-vectors --index-name my-index --namespace my-namespace --all-vectors + pc index delete-vectors --index-name my-index --namespace my-namespace --filter '{"genre": "classical"}' + `), + Run: func(cmd *cobra.Command, args []string) { + runDeleteVectorsCmd(cmd.Context(), options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of the index to delete vectors from") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to delete vectors from") + cmd.Flags().Var(&options.ids, "ids", "IDs of the vectors to delete") + cmd.Flags().Var(&options.filter, "filter", "filter to delete the vectors with") + cmd.Flags().BoolVar(&options.deleteAllVectors, "all-vectors", false, "delete all vectors from the namespace") + cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON") + + _ = cmd.MarkFlagRequired("index-name") + + return cmd +} + +func runDeleteVectorsCmd(ctx context.Context, options deleteVectorsCmdOptions) { + pc := sdk.NewPineconeClient(ctx) + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + if options.ids == nil && options.filter == nil && !options.deleteAllVectors { + msg.FailMsg("Either --ids, --filter, or --all-vectors must be provided") + exit.ErrorMsg("Either --ids, --filter, or --all-vectors must be provided") + } + + // Delete all vectors in namespace + if options.deleteAllVectors { + err = ic.DeleteAllVectorsInNamespace(ctx) + if err != nil { + msg.FailMsg("Failed to delete all vectors in namespace: %s", err) + exit.Error(err, "Failed to delete all vectors in namespace") + } + if !options.json { + msg.SuccessMsg("Deleted all vectors in namespace: %s", options.namespace) + } + return + } + + // Delete vectors by ID + if len(options.ids) > 0 { + err = ic.DeleteVectorsById(ctx, options.ids) + if err != nil { + msg.FailMsg("Failed to delete vectors by IDs: %s", err) + exit.Error(err, "Failed to delete vectors by IDs") + } + if !options.json { + msg.SuccessMsg("Deleted vectors by IDs: %s", options.ids) + } + return + } + + // Delete vectors by filter + if options.filter != nil { + filter, err := pinecone.NewMetadataFilter(options.filter) + if err != nil { + msg.FailMsg("Failed to create filter: %s", err) + exit.Errorf(err, "Failed to create filter") + } + + err = ic.DeleteVectorsByFilter(ctx, filter) + if err != nil { + msg.FailMsg("Failed to delete vectors by filter: %s", err) + exit.Error(err, "Failed to delete vectors by filter") + } + if !options.json { + msg.SuccessMsg("Deleted vectors by filter: %s", filter.String()) + } + return + } +} diff --git a/internal/pkg/cli/command/index/describe.go b/internal/pkg/cli/command/index/describe.go index 93f8165..1588c07 100644 --- a/internal/pkg/cli/command/index/describe.go +++ b/internal/pkg/cli/command/index/describe.go @@ -29,7 +29,8 @@ func NewDescribeCmd() *cobra.Command { pc index describe --name "index-name" `), Run: func(cmd *cobra.Command, args []string) { - pc := sdk.NewPineconeClient() + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) idx, err := pc.DescribeIndex(cmd.Context(), options.name) if err != nil { diff --git a/internal/pkg/cli/command/index/describe_stats.go b/internal/pkg/cli/command/index/describe_stats.go new file mode 100644 index 0000000..e818d42 --- /dev/null +++ b/internal/pkg/cli/command/index/describe_stats.go @@ -0,0 +1,78 @@ +package index + +import ( + "context" + + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/flags" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/presenters" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" + "github.com/spf13/cobra" +) + +type describeStatsCmdOptions struct { + indexName string + filter flags.JSONObject + json bool +} + +func NewDescribeIndexStatsCmd() *cobra.Command { + options := describeStatsCmdOptions{} + cmd := &cobra.Command{ + Use: "describe-stats", + Short: "Describe the stats of an index", + Example: help.Examples(` + pc index describe-stats --index-name "index-name" + pc index describe-stats --index-name "index-name" --filter '{"k":"v"}' + pc index describe-stats --index-name "index-name" --filter @./filter.json + `), + Run: func(cmd *cobra.Command, args []string) { + runDescribeIndexStatsCmd(cmd.Context(), options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of index to describe stats for") + cmd.Flags().VarP(&options.filter, "filter", "f", "metadata filter to apply to the operation") + cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON") + _ = cmd.MarkFlagRequired("index-name") + + return cmd +} + +func runDescribeIndexStatsCmd(ctx context.Context, options describeStatsCmdOptions) { + pc := sdk.NewPineconeClient(ctx) + + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, "") + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + // Build metadata filter if provided + var filter *pinecone.MetadataFilter + if options.filter != nil { + filter, err = pinecone.NewMetadataFilter(options.filter) + if err != nil { + msg.FailMsg("Failed to create filter: %s", err) + exit.Errorf(err, "Failed to create filter") + } + } + + resp, err := ic.DescribeIndexStatsFiltered(ctx, filter) + if err != nil { + msg.FailMsg("Failed to describe stats: %s", err) + exit.Error(err, "Failed to describe stats") + } + + if options.json { + json := text.IndentJSON(resp) + pcio.Println(json) + } else { + presenters.PrintDescribeIndexStatsTable(resp) + } +} diff --git a/internal/pkg/cli/command/index/fetch.go b/internal/pkg/cli/command/index/fetch.go new file mode 100644 index 0000000..8c3fa80 --- /dev/null +++ b/internal/pkg/cli/command/index/fetch.go @@ -0,0 +1,156 @@ +package index + +import ( + "context" + + "github.com/pinecone-io/cli/internal/pkg/utils/bodyutil" + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/flags" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/presenters" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/style" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" + "github.com/spf13/cobra" +) + +type fetchBody struct { + Ids []string `json:"ids"` + Filter map[string]any `json:"filter"` + Limit *uint32 `json:"limit"` + PaginationToken *string `json:"pagination_token"` +} + +type fetchCmdOptions struct { + indexName string + namespace string + ids flags.StringList + filter flags.JSONObject + limit uint32 + paginationToken string + body string + json bool +} + +func NewFetchCmd() *cobra.Command { + options := fetchCmdOptions{} + cmd := &cobra.Command{ + Use: "fetch", + Short: "Fetch vectors by ID or metadata filter from an index", + Example: help.Examples(` + pc index fetch --index-name my-index --ids '["123","456","789"]' + pc index fetch --index-name my-index --ids @./ids.json + + pc index fetch --index-name my-index --filter '{"key": "value"}' + pc index fetch --index-name my-index --filter @./filter.json + + pc index fetch --index-name my-index --body @./fetch.json + cat fetch.json | pc index fetch --index-name my-index --body @- + `), + Run: func(cmd *cobra.Command, args []string) { + runFetchCmd(cmd.Context(), options) + }, + } + + cmd.Flags().VarP(&options.ids, "ids", "i", "IDs of vectors to fetch (inline JSON array, @path.json, or @- for stdin)") + cmd.Flags().VarP(&options.filter, "filter", "f", "metadata filter to apply to the fetch (inline JSON, @path.json, or @- for stdin)") + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of the index to fetch from") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to fetch from") + cmd.Flags().Uint32VarP(&options.limit, "limit", "l", 0, "maximum number of vectors to fetch") + cmd.Flags().StringVarP(&options.paginationToken, "pagination-token", "p", "", "pagination token to continue a previous listing operation") + cmd.Flags().StringVar(&options.body, "body", "", "request body JSON (inline, @path.json, or @- for stdin; only one argument may use stdin)") + cmd.Flags().BoolVarP(&options.json, "json", "j", false, "output as JSON") + + cmd.MarkFlagsMutuallyExclusive("ids", "filter") + _ = cmd.MarkFlagRequired("index-name") + + return cmd +} + +func runFetchCmd(ctx context.Context, options fetchCmdOptions) { + pc := sdk.NewPineconeClient(ctx) + + // Apply body overlay if provided + if options.body != "" { + if b, src, err := bodyutil.DecodeBodyArgs[fetchBody](options.body); err != nil { + msg.FailMsg("Failed to parse fetch body (%s): %s", style.Emphasis(src.Label), err) + exit.Errorf(err, "Failed to parse fetch body (%s): %v", src.Label, err) + } else if b != nil { + if len(options.ids) == 0 && len(b.Ids) > 0 { + options.ids = b.Ids + } + if options.filter == nil && b.Filter != nil { + options.filter = b.Filter + } + if b.Limit != nil { + options.limit = *b.Limit + } + if b.PaginationToken != nil { + options.paginationToken = *b.PaginationToken + } + } + } + + if len(options.ids) > 0 && (options.limit > 0 || options.paginationToken != "") { + msg.FailMsg("ids and limit/pagination-token cannot be used together") + exit.ErrorMsg("ids and limit/pagination-token cannot be used together") + } + + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + if options.ids == nil && options.filter == nil { + msg.FailMsg("Either --ids or --filter must be provided") + exit.ErrorMsg("Either --ids or --filter must be provided") + } + + // Fetch vectors by ID + if len(options.ids) > 0 { + vectors, err := ic.FetchVectors(ctx, []string(options.ids)) + if err != nil { + exit.Error(err, "Failed to fetch vectors") + } + printFetchVectorsResults(presenters.NewFetchVectorsResultsFromFetch(vectors), options) + } + + // Fetch vectors by metadata filter + if options.filter != nil { + filter, err := pinecone.NewMetadataFilter(options.filter) + if err != nil { + msg.FailMsg("Failed to create filter: %s", err) + exit.Errorf(err, "Failed to create filter") + } + + req := &pinecone.FetchVectorsByMetadataRequest{ + Filter: filter, + } + + if options.limit > 0 { + req.Limit = &options.limit + } + if options.paginationToken != "" { + req.PaginationToken = &options.paginationToken + } + + vectors, err := ic.FetchVectorsByMetadata(ctx, req) + if err != nil { + exit.Error(err, "Failed to fetch vectors by metadata") + } + printFetchVectorsResults(presenters.NewFetchVectorsResultsFromFetchByMetadata(vectors), options) + } +} + +func printFetchVectorsResults(results *presenters.FetchVectorsResults, options fetchCmdOptions) { + if options.json { + json := text.IndentJSON(results) + pcio.Println(json) + } else { + presenters.PrintFetchVectorsTable(results) + } +} diff --git a/internal/pkg/cli/command/index/list.go b/internal/pkg/cli/command/index/list.go index 5e632df..4a7a858 100644 --- a/internal/pkg/cli/command/index/list.go +++ b/internal/pkg/cli/command/index/list.go @@ -1,7 +1,6 @@ package index import ( - "context" "os" "sort" "strings" @@ -32,8 +31,8 @@ func NewListCmd() *cobra.Command { pc index list `), Run: func(cmd *cobra.Command, args []string) { - pc := sdk.NewPineconeClient() - ctx := context.Background() + ctx := cmd.Context() + pc := sdk.NewPineconeClient(ctx) idxs, err := pc.ListIndexes(ctx) if err != nil { diff --git a/internal/pkg/cli/command/index/list_vectors.go b/internal/pkg/cli/command/index/list_vectors.go new file mode 100644 index 0000000..099e522 --- /dev/null +++ b/internal/pkg/cli/command/index/list_vectors.go @@ -0,0 +1,83 @@ +package index + +import ( + "context" + + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/presenters" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" + "github.com/spf13/cobra" +) + +type listVectorsCmdOptions struct { + indexName string + namespace string + limit uint32 + paginationToken string + json bool +} + +func NewListVectorsCmd() *cobra.Command { + options := listVectorsCmdOptions{} + cmd := &cobra.Command{ + Use: "list-vectors", + Short: "List vectors in an index", + Example: help.Examples(` + pc index list-vectors --index-name my-index --namespace my-namespace + `), + Run: func(cmd *cobra.Command, args []string) { + runListVectorsCmd(cmd.Context(), options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of the index to list vectors from") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to list vectors from") + cmd.Flags().Uint32VarP(&options.limit, "limit", "l", 0, "maximum number of vectors to list") + cmd.Flags().StringVarP(&options.paginationToken, "pagination-token", "p", "", "pagination token to continue a previous listing operation") + cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON") + + _ = cmd.MarkFlagRequired("index-name") + + return cmd +} + +func runListVectorsCmd(ctx context.Context, options listVectorsCmdOptions) { + pc := sdk.NewPineconeClient(ctx) + + // Get IndexConnection + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + var limit *uint32 + if options.limit > 0 { + limit = &options.limit + } + var paginationToken *string + if options.paginationToken != "" { + paginationToken = &options.paginationToken + } + + resp, err := ic.ListVectors(ctx, &pinecone.ListVectorsRequest{ + Limit: limit, + PaginationToken: paginationToken, + }) + if err != nil { + msg.FailMsg("Failed to list vectors: %s", err) + exit.Error(err, "Failed to list vectors") + } + + if options.json { + json := text.IndentJSON(resp) + pcio.Println(json) + } else { + presenters.PrintListVectorsTable(resp) + } +} diff --git a/internal/pkg/cli/command/index/query.go b/internal/pkg/cli/command/index/query.go new file mode 100644 index 0000000..d2d9ad8 --- /dev/null +++ b/internal/pkg/cli/command/index/query.go @@ -0,0 +1,201 @@ +package index + +import ( + "context" + + "github.com/pinecone-io/cli/internal/pkg/utils/bodyutil" + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/flags" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/presenters" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/style" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" + "github.com/spf13/cobra" +) + +type queryBody struct { + Id string `json:"id"` + Vector []float32 `json:"vector"` + SparseValues *pinecone.SparseValues `json:"sparse_values"` + Filter map[string]any `json:"filter"` + TopK *uint32 `json:"top_k"` + IncludeValues *bool `json:"include_values"` + IncludeMetadata *bool `json:"include_metadata"` +} + +type queryCmdOptions struct { + id string + vector flags.Float32List + sparseIndices flags.UInt32List + sparseValues flags.Float32List + indexName string + namespace string + topK uint32 + filter flags.JSONObject + includeValues bool + includeMetadata bool + body string + json bool +} + +func NewQueryCmd() *cobra.Command { + options := queryCmdOptions{} + cmd := &cobra.Command{ + Use: "query", + Short: "Query an index by vector values", + Example: help.Examples(` + pc index query --index-name my-index --id doc-123 --top-k 10 --include-metadata + + pc index query --index-name my-index --vector '[0.1, 0.2, 0.3]' --top-k 25 + pc index query --index-name my-index --vector @./vector.json --top-k 25 --include-metadata + jq -c '.embedding' doc.json | pc index query --index-name my-index --vector @- --top-k 20 + + pc index query --index-name my-index --sparse-indices @./indices.json --sparse-values @./values.json --top-k 15 + + pc index query --index-name my-index --vector @./vector.json --filter '{"genre":"sci-fi"}' --include-metadata + + pc index query --index-name my-index --body @./query.json + cat query.json | pc index query --index-name my-index --body @- + `), + Run: func(cmd *cobra.Command, args []string) { + runQueryCmd(cmd.Context(), options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of the index to query") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "index namespace to query") + cmd.Flags().Uint32VarP(&options.topK, "top-k", "k", 10, "maximum number of results to return") + cmd.Flags().VarP(&options.filter, "filter", "f", "metadata filter to apply to the query (inline JSON, @path.json, or @- for stdin)") + cmd.Flags().BoolVar(&options.includeValues, "include-values", false, "include vector values in the query results") + cmd.Flags().BoolVar(&options.includeMetadata, "include-metadata", false, "include metadata in the query results") + cmd.Flags().StringVarP(&options.id, "id", "i", "", "ID of the vector to query against") + cmd.Flags().VarP(&options.vector, "vector", "v", "vector values to query against (inline JSON array, @path.json, or @- for stdin)") + cmd.Flags().Var(&options.sparseIndices, "sparse-indices", "sparse indices to query against (inline JSON array, @path.json, or @- for stdin)") + cmd.Flags().Var(&options.sparseValues, "sparse-values", "sparse values to query against (inline JSON array, @path.json, or @- for stdin)") + cmd.Flags().StringVar(&options.body, "body", "", "request body JSON (inline, @path.json, or @- for stdin; only one argument may use stdin)") + cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON") + + _ = cmd.MarkFlagRequired("index-name") + cmd.MarkFlagsMutuallyExclusive("id", "vector", "sparse-values") + + return cmd +} + +func runQueryCmd(ctx context.Context, options queryCmdOptions) { + pc := sdk.NewPineconeClient(ctx) + + // Apply body overlay if provided + if options.body != "" { + if b, src, err := bodyutil.DecodeBodyArgs[queryBody](options.body); err != nil { + msg.FailMsg("Failed to parse query body (%s): %s", style.Emphasis(src.Label), err) + exit.Errorf(err, "Failed to parse query body (%s): %v", src.Label, err) + } else if b != nil { + if options.id == "" && b.Id != "" { + options.id = b.Id + } + if len(options.vector) == 0 && len(b.Vector) > 0 { + options.vector = b.Vector + } + if (len(options.sparseIndices) == 0 && len(options.sparseValues) == 0) && b.SparseValues != nil { + options.sparseIndices = b.SparseValues.Indices + options.sparseValues = b.SparseValues.Values + } + if options.filter == nil && b.Filter != nil { + options.filter = b.Filter + } + if b.TopK != nil { + options.topK = *b.TopK + } + if b.IncludeValues != nil { + options.includeValues = *b.IncludeValues + } + if b.IncludeMetadata != nil { + options.includeMetadata = *b.IncludeMetadata + } + } + } + + if options.id == "" && options.vector == nil && options.sparseIndices == nil && options.sparseValues == nil && options.filter == nil { + msg.FailMsg("Either --id, --vector, --sparse-indices & --sparse-values, or --filter must be provided") + exit.ErrorMsg("Either --id, --vector, --sparse-indices & --sparse-values, or --filter must be provided") + } + + // Get IndexConnection + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + var queryResponse *pinecone.QueryVectorsResponse + + // Build metadata filter if provided + var filter *pinecone.MetadataFilter + if options.filter != nil { + filter, err = pinecone.NewMetadataFilter(options.filter) + if err != nil { + msg.FailMsg("Failed to create filter: %s", err) + exit.Errorf(err, "Failed to create filter") + } + } + + // Query by vector ID + if options.id != "" { + req := &pinecone.QueryByVectorIdRequest{ + VectorId: options.id, + TopK: options.topK, + IncludeValues: options.includeValues, + IncludeMetadata: options.includeMetadata, + MetadataFilter: filter, + } + + queryResponse, err = ic.QueryByVectorId(ctx, req) + if err != nil { + exit.Error(err, "Failed to query by vector ID") + } + } + + // Query by vector values + if len(options.vector) > 0 || len(options.sparseIndices) > 0 || len(options.sparseValues) > 0 { + var sparse *pinecone.SparseValues + + // Only include sparse values if the user provided them + if len(options.sparseIndices) > 0 || len(options.sparseValues) > 0 { + if len(options.sparseIndices) == 0 || len(options.sparseValues) == 0 { + exit.Errorf(nil, "both --sparse-indices and --sparse-values are required when specifying sparse values") + } + if len(options.sparseIndices) != len(options.sparseValues) { + exit.Errorf(nil, "--sparse-indices and --sparse-values must be the same length") + } + sparse = &pinecone.SparseValues{ + Indices: options.sparseIndices, + Values: options.sparseValues, + } + } + + req := &pinecone.QueryByVectorValuesRequest{ + Vector: options.vector, + SparseValues: sparse, + TopK: options.topK, + IncludeValues: options.includeValues, + IncludeMetadata: options.includeMetadata, + MetadataFilter: filter, + } + + queryResponse, err = ic.QueryByVectorValues(ctx, req) + if err != nil { + exit.Error(err, "Failed to query by vector values") + } + } + + if options.json { + json := text.IndentJSON(queryResponse) + pcio.Println(json) + } else { + presenters.PrintQueryVectorsTable(queryResponse) + } +} diff --git a/internal/pkg/cli/command/index/update.go b/internal/pkg/cli/command/index/update.go new file mode 100644 index 0000000..16b3a97 --- /dev/null +++ b/internal/pkg/cli/command/index/update.go @@ -0,0 +1,194 @@ +package index + +import ( + "context" + + "github.com/pinecone-io/cli/internal/pkg/utils/bodyutil" + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/flags" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/presenters" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/style" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" + "github.com/spf13/cobra" +) + +type updateBody struct { + Id string `json:"id"` + Values []float32 `json:"values"` + SparseValues *pinecone.SparseValues `json:"sparse_values"` + Metadata map[string]any `json:"metadata"` + Filter map[string]any `json:"filter"` + DryRun *bool `json:"dry_run"` +} + +type updateCmdOptions struct { + indexName string + namespace string + id string + values flags.Float32List + sparseIndices flags.UInt32List + sparseValues flags.Float32List + metadata flags.JSONObject + filter flags.JSONObject + dryRun bool + body string + json bool +} + +func NewUpdateCmd() *cobra.Command { + options := updateCmdOptions{} + + cmd := &cobra.Command{ + Use: "update", + Short: "Update a vector by ID, or a set of vectors by metadata filter", + Example: help.Examples(` + pc index update --index-name my-index --id doc-123 --values '[0.1, 0.2, 0.3]' + pc index update --index-name my-index --id doc-123 --sparse-indices @./indices.json --sparse-values @./values.json + pc index update --index-name my-index --id doc-123 --metadata '{"genre":"sci-fi"}' + pc index update --index-name my-index --filter '{"genre":"sci-fi"}' --metadata '{"genre":"fantasy"}' --dry-run + pc index update --index-name my-index --body @./update.json + cat update.json | pc index update --index-name my-index --body @- + `), + Run: func(cmd *cobra.Command, args []string) { + runUpdateCmd(cmd.Context(), options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of the index to update") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to update the vector in") + cmd.Flags().StringVar(&options.id, "id", "", "ID of the vector to update") + cmd.Flags().Var(&options.values, "values", "values to update the vector with (inline JSON array, @path.json, or @- for stdin)") + cmd.Flags().Var(&options.sparseIndices, "sparse-indices", "sparse indices to update the vector with (inline JSON array, @path.json, or @- for stdin)") + cmd.Flags().Var(&options.sparseValues, "sparse-values", "sparse values to update the vector with (inline JSON array, @path.json, or @- for stdin)") + cmd.Flags().Var(&options.metadata, "metadata", "metadata to update the vector with (inline JSON, @path.json, or @- for stdin)") + cmd.Flags().Var(&options.filter, "filter", "filter to update the vectors with (inline JSON, @path.json, or @- for stdin)") + cmd.Flags().BoolVar(&options.dryRun, "dry-run", false, "do not update the vectors, just return the number of vectors that would be updated") + cmd.Flags().StringVar(&options.body, "body", "", "request body JSON (inline, @path.json, or @- for stdin; only one argument may use stdin)") + cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON") + + _ = cmd.MarkFlagRequired("index-name") + cmd.MarkFlagsMutuallyExclusive("id", "filter") + + return cmd +} + +func runUpdateCmd(ctx context.Context, options updateCmdOptions) { + pc := sdk.NewPineconeClient(ctx) + + // Apply body overlay if provided + if options.body != "" { + if b, src, err := bodyutil.DecodeBodyArgs[updateBody](options.body); err != nil { + msg.FailMsg("Failed to parse update body (%s): %s", style.Emphasis(src.Label), err) + exit.Errorf(err, "Failed to parse update body (%s): %v", src.Label, err) + } else if b != nil { + if options.id == "" && b.Id != "" { + options.id = b.Id + } + if len(options.values) == 0 && len(b.Values) > 0 { + options.values = b.Values + } + if (len(options.sparseIndices) == 0 && len(options.sparseValues) == 0) && b.SparseValues != nil { + options.sparseIndices = b.SparseValues.Indices + options.sparseValues = b.SparseValues.Values + } + if options.filter == nil && b.Filter != nil { + options.filter = b.Filter + } + if options.metadata == nil && b.Metadata != nil { + options.metadata = b.Metadata + } + if b.DryRun != nil { + options.dryRun = *b.DryRun + } + } + } + + // Validate update by ID or metadata filter + if options.id == "" && options.filter == nil { + msg.FailMsg("Either --id or --filter must be provided") + exit.ErrorMsg("Either --id or --filter must be provided") + } + + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + // Update vector by ID + if options.id != "" { + metadata, err := pinecone.NewMetadata(options.metadata) + if err != nil { + msg.FailMsg("Failed to create metadata: %s", err) + exit.Errorf(err, "Failed to create metadata") + } + + var sparseValues *pinecone.SparseValues + if len(options.sparseIndices) > 0 || len(options.sparseValues) > 0 { + sparseValues = &pinecone.SparseValues{ + Indices: options.sparseIndices, + Values: options.sparseValues, + } + } + + err = ic.UpdateVector(ctx, &pinecone.UpdateVectorRequest{ + Id: options.id, + Values: options.values, + SparseValues: sparseValues, + Metadata: metadata, + }) + if err != nil { + msg.FailMsg("Failed to update vector ID: %s - %v", options.id, err) + exit.Errorf(err, "Failed to update vector ID: %s", options.id) + } + + if !options.json { + msg.SuccessMsg("Vector ID: %s updated successfully", options.id) + } + return + } + + // Update vectors by metadata filter + if options.filter != nil { + filter, err := pinecone.NewMetadataFilter(options.filter) + if err != nil { + msg.FailMsg("Failed to create filter: %s", err) + exit.Errorf(err, "Failed to create filter") + } + + metadata, err := pinecone.NewMetadata(options.metadata) + if err != nil { + msg.FailMsg("Failed to create metadata: %s", err) + exit.Errorf(err, "Failed to create metadata") + } + + var dryRun *bool + if options.dryRun { + dryRun = &options.dryRun + } + + resp, err := ic.UpdateVectorsByMetadata(ctx, &pinecone.UpdateVectorsByMetadataRequest{ + Filter: filter, + Metadata: metadata, + DryRun: dryRun, + }) + if err != nil { + msg.FailMsg("Failed to update vectors by metadata: %s - %v", filter.String(), err) + exit.Errorf(err, "Failed to update vectors by metadata: %s", filter.String()) + } + + if !options.json { + msg.SuccessMsg("Updated %d vectors by metadata filter: %s", resp.MatchedRecords, filter.String()) + presenters.PrintUpdateVectorsByMetadataTable(resp) + } else { + json := text.IndentJSON(resp) + pcio.Println(json) + } + return + } +} diff --git a/internal/pkg/cli/command/index/upsert.go b/internal/pkg/cli/command/index/upsert.go new file mode 100644 index 0000000..f3e958d --- /dev/null +++ b/internal/pkg/cli/command/index/upsert.go @@ -0,0 +1,126 @@ +package index + +import ( + "context" + + "github.com/pinecone-io/cli/internal/pkg/utils/bodyutil" + "github.com/pinecone-io/cli/internal/pkg/utils/exit" + "github.com/pinecone-io/cli/internal/pkg/utils/help" + "github.com/pinecone-io/cli/internal/pkg/utils/msg" + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/sdk" + "github.com/pinecone-io/cli/internal/pkg/utils/style" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/spf13/cobra" + + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +type upsertBody struct { + Vectors []pinecone.Vector `json:"vectors"` +} + +type upsertCmdOptions struct { + body string + indexName string + namespace string + batchSize int + json bool +} + +func NewUpsertCmd() *cobra.Command { + options := upsertCmdOptions{} + + cmd := &cobra.Command{ + Use: "upsert [file]", + Short: "Upsert vectors into an index from a JSON file", + Example: help.Examples(` + pc index upsert --index-name my-index --namespace my-namespace ./vectors.json + pc index upsert --index-name my-index --namespace my-namespace --file - < ./vectors.json + pc index upsert --index-name my-index --body @./payload.json + cat payload.json | pc index upsert --index-name my-index --body @- + `), + Run: func(cmd *cobra.Command, args []string) { + runUpsertCmd(cmd.Context(), options) + }, + } + + cmd.Flags().StringVarP(&options.indexName, "index-name", "n", "", "name of index to upsert into") + cmd.Flags().StringVar(&options.namespace, "namespace", "__default__", "namespace to upsert into") + cmd.Flags().StringVar(&options.body, "body", "", "request body JSON (inline, @path.json, or @- for stdin; only one argument may use stdin)") + cmd.Flags().IntVarP(&options.batchSize, "batch-size", "b", 1000, "size of batches to upsert (default: 1000)") + cmd.Flags().BoolVar(&options.json, "json", false, "output as JSON") + _ = cmd.MarkFlagRequired("index-name") + _ = cmd.MarkFlagRequired("body") + + return cmd +} + +func runUpsertCmd(ctx context.Context, options upsertCmdOptions) { + var payload *upsertBody + payload, src, err := bodyutil.DecodeBodyArgs[upsertBody](options.body) + if err != nil { + msg.FailMsg("Failed to parse upsert body (%s): %s", style.Emphasis(src.Label), err) + exit.Error(err, "Failed to parse upsert body") + } + + // Get IndexConnection + pc := sdk.NewPineconeClient(ctx) + ic, err := sdk.NewIndexConnection(ctx, pc, options.indexName, options.namespace) + if err != nil { + msg.FailMsg("Failed to create index connection: %s", err) + exit.Error(err, "Failed to create index connection") + } + + if len(payload.Vectors) == 0 { + msg.FailMsg("No vectors found in %s", style.Emphasis(src.Label)) + exit.ErrorMsg("No vectors provided for upsert") + } + + // Map to SDK types + mapped := make([]*pinecone.Vector, 0, len(payload.Vectors)) + for _, v := range payload.Vectors { + values := v.Values + + var vector pinecone.Vector + vector.Id = v.Id + vector.Metadata = v.Metadata + + if values != nil { + vector.Values = values + } + + if v.SparseValues != nil { + vector.SparseValues = &pinecone.SparseValues{ + Indices: v.SparseValues.Indices, + Values: v.SparseValues.Values, + } + } + + mapped = append(mapped, &vector) + } + + batches := make([][]*pinecone.Vector, 0, (len(mapped)+options.batchSize-1)/options.batchSize) + for i := 0; i < len(mapped); i += options.batchSize { + end := i + options.batchSize + if end > len(mapped) { + end = len(mapped) + } + batches = append(batches, mapped[i:end]) + } + + for i, batch := range batches { + resp, err := ic.UpsertVectors(ctx, batch) + if err != nil { + msg.FailMsg("Failed to upsert %d vectors in batch %d: %s", len(batch), i+1, err) + exit.Errorf(err, "Failed to upsert %d vectors in batch %d", len(batch), i+1) + } else { + if options.json { + json := text.IndentJSON(resp) + pcio.Println(json) + } else { + msg.SuccessMsg("Upserted %d vectors into namespace %s (batch %d of %d)", len(batch), options.namespace, i+1, len(batches)) + } + } + } +} diff --git a/internal/pkg/cli/command/organization/delete.go b/internal/pkg/cli/command/organization/delete.go index b26471b..e0f9f81 100644 --- a/internal/pkg/cli/command/organization/delete.go +++ b/internal/pkg/cli/command/organization/delete.go @@ -34,7 +34,8 @@ func NewDeleteOrganizationCmd() *cobra.Command { `), GroupID: help.GROUP_ORGANIZATIONS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) // get the organization first org, err := ac.Organization.Describe(cmd.Context(), options.organizationID) diff --git a/internal/pkg/cli/command/organization/describe.go b/internal/pkg/cli/command/organization/describe.go index 833b0d0..d6ae4f6 100644 --- a/internal/pkg/cli/command/organization/describe.go +++ b/internal/pkg/cli/command/organization/describe.go @@ -29,7 +29,8 @@ func NewDescribeOrganizationCmd() *cobra.Command { `), GroupID: help.GROUP_ORGANIZATIONS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) orgId := options.organizationID var err error diff --git a/internal/pkg/cli/command/organization/list.go b/internal/pkg/cli/command/organization/list.go index 9feea27..91580f6 100644 --- a/internal/pkg/cli/command/organization/list.go +++ b/internal/pkg/cli/command/organization/list.go @@ -31,7 +31,8 @@ func NewListOrganizationsCmd() *cobra.Command { `), GroupID: help.GROUP_ORGANIZATIONS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) orgs, err := ac.Organization.List(cmd.Context()) if err != nil { diff --git a/internal/pkg/cli/command/organization/update.go b/internal/pkg/cli/command/organization/update.go index fedae0e..5549791 100644 --- a/internal/pkg/cli/command/organization/update.go +++ b/internal/pkg/cli/command/organization/update.go @@ -31,7 +31,8 @@ func NewUpdateOrganizationCmd() *cobra.Command { `), GroupID: help.GROUP_ORGANIZATIONS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) orgId := options.organizationID var err error diff --git a/internal/pkg/cli/command/project/create.go b/internal/pkg/cli/command/project/create.go index e8c422f..1359750 100644 --- a/internal/pkg/cli/command/project/create.go +++ b/internal/pkg/cli/command/project/create.go @@ -1,8 +1,6 @@ package project import ( - "context" - "github.com/pinecone-io/cli/internal/pkg/utils/configuration/state" "github.com/pinecone-io/cli/internal/pkg/utils/exit" "github.com/pinecone-io/cli/internal/pkg/utils/help" @@ -49,7 +47,8 @@ func NewCreateProjectCmd() *cobra.Command { pc project create --name "demo-project" --max-pods 10 --force-encryption `), Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) createParams := &pinecone.CreateProjectParams{} if options.name != "" { @@ -62,7 +61,7 @@ func NewCreateProjectCmd() *cobra.Command { createParams.ForceEncryptionWithCmek = &options.forceEncryptionWithCMEK } - proj, err := ac.Project.Create(context.Background(), createParams) + proj, err := ac.Project.Create(ctx, createParams) if err != nil { msg.FailMsg("Failed to create project %s: %s\n", style.Emphasis(options.name), err) exit.Errorf(err, "Failed to create project %s", style.Emphasis(options.name)) diff --git a/internal/pkg/cli/command/project/delete.go b/internal/pkg/cli/command/project/delete.go index 69dbd42..38e537e 100644 --- a/internal/pkg/cli/command/project/delete.go +++ b/internal/pkg/cli/command/project/delete.go @@ -38,8 +38,8 @@ func NewDeleteProjectCmd() *cobra.Command { `), GroupID: help.GROUP_PROJECTS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() - ctx := context.Background() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) projId := options.projectId var err error @@ -58,8 +58,8 @@ func NewDeleteProjectCmd() *cobra.Command { exit.Error(err, "Failed to retrieve project information") } - verifyNoIndexes(projToDelete.Id, projToDelete.Name) - verifyNoCollections(projToDelete.Id, projToDelete.Name) + verifyNoIndexes(ctx, projToDelete.Id, projToDelete.Name) + verifyNoCollections(ctx, projToDelete.Id, projToDelete.Name) if !options.skipConfirmation { confirmDelete(projToDelete.Name) @@ -117,10 +117,9 @@ func confirmDelete(projectName string) { } } -func verifyNoIndexes(projectId string, projectName string) { +func verifyNoIndexes(ctx context.Context, projectId string, projectName string) { // Check if project contains indexes - pc := sdk.NewPineconeClientForProjectById(projectId) - ctx := context.Background() + pc := sdk.NewPineconeClientForProjectById(ctx, projectId) idxs, err := pc.ListIndexes(ctx) if err != nil { @@ -134,10 +133,9 @@ func verifyNoIndexes(projectId string, projectName string) { } } -func verifyNoCollections(projectId string, projectName string) { +func verifyNoCollections(ctx context.Context, projectId string, projectName string) { // Check if project contains collections - pc := sdk.NewPineconeClientForProjectById(projectId) - ctx := context.Background() + pc := sdk.NewPineconeClientForProjectById(ctx, projectId) collections, err := pc.ListCollections(ctx) if err != nil { diff --git a/internal/pkg/cli/command/project/describe.go b/internal/pkg/cli/command/project/describe.go index edf85b7..9558b7e 100644 --- a/internal/pkg/cli/command/project/describe.go +++ b/internal/pkg/cli/command/project/describe.go @@ -1,8 +1,6 @@ package project import ( - "context" - "github.com/pinecone-io/cli/internal/pkg/utils/configuration/state" "github.com/pinecone-io/cli/internal/pkg/utils/exit" "github.com/pinecone-io/cli/internal/pkg/utils/help" @@ -31,7 +29,8 @@ func NewDescribeProjectCmd() *cobra.Command { pc project describe --id "project-id" `), Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) projId := options.projectID var err error @@ -43,7 +42,7 @@ func NewDescribeProjectCmd() *cobra.Command { } } - project, err := ac.Project.Describe(context.Background(), projId) + project, err := ac.Project.Describe(ctx, projId) if err != nil { msg.FailMsg("Failed to describe project %s: %s\n", projId, err) exit.Errorf(err, "Failed to describe project %s", style.Emphasis(projId)) diff --git a/internal/pkg/cli/command/project/list.go b/internal/pkg/cli/command/project/list.go index 3bae081..8014924 100644 --- a/internal/pkg/cli/command/project/list.go +++ b/internal/pkg/cli/command/project/list.go @@ -1,7 +1,6 @@ package project import ( - "context" "os" "strconv" "strings" @@ -33,8 +32,8 @@ func NewListProjectsCmd() *cobra.Command { pc project list `), Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() - ctx := context.Background() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) projects, err := ac.Project.List(ctx) if err != nil { diff --git a/internal/pkg/cli/command/project/update.go b/internal/pkg/cli/command/project/update.go index 10ef925..1d40d36 100644 --- a/internal/pkg/cli/command/project/update.go +++ b/internal/pkg/cli/command/project/update.go @@ -1,8 +1,6 @@ package project import ( - "context" - "github.com/pinecone-io/cli/internal/pkg/utils/configuration/state" "github.com/pinecone-io/cli/internal/pkg/utils/exit" "github.com/pinecone-io/cli/internal/pkg/utils/help" @@ -35,7 +33,8 @@ func NewUpdateProjectCmd() *cobra.Command { `), GroupID: help.GROUP_PROJECTS.ID, Run: func(cmd *cobra.Command, args []string) { - ac := sdk.NewPineconeAdminClient() + ctx := cmd.Context() + ac := sdk.NewPineconeAdminClient(ctx) projId := options.projectId var err error @@ -60,7 +59,7 @@ func NewUpdateProjectCmd() *cobra.Command { updateParams.MaxPods = &options.maxPods } - project, err := ac.Project.Update(context.Background(), projId, updateParams) + project, err := ac.Project.Update(ctx, projId, updateParams) if err != nil { msg.FailMsg("Failed to update project %s: %s\n", projId, err) exit.Errorf(err, "Failed to update project %s", style.Emphasis(projId)) diff --git a/internal/pkg/cli/command/root/root.go b/internal/pkg/cli/command/root/root.go index 8092db3..290d0cf 100644 --- a/internal/pkg/cli/command/root/root.go +++ b/internal/pkg/cli/command/root/root.go @@ -1,7 +1,11 @@ package root import ( + "context" "os" + "os/signal" + "syscall" + "time" "github.com/pinecone-io/cli/internal/pkg/cli/command/apiKey" "github.com/pinecone-io/cli/internal/pkg/cli/command/auth" @@ -20,14 +24,23 @@ import ( "github.com/spf13/cobra" ) -var rootCmd *cobra.Command +var ( + rootCmd *cobra.Command + globalOptions GlobalOptions + cancelRootFunc context.CancelFunc +) type GlobalOptions struct { - quiet bool + quiet bool + timeout time.Duration } func Execute() { - err := rootCmd.Execute() + //Base context: cancel on SIGINT / SIGTERM + ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer cancel() + + err := rootCmd.ExecuteContext(ctx) if err != nil { os.Exit(1) } @@ -62,12 +75,30 @@ var ( ) func init() { - globalOptions := GlobalOptions{} + // Default timeout for context.Context cancellation + // This is applied to individual operations within subcommands through cmd.Context() + defaultTimeout := 60 * time.Second + globalOptions = GlobalOptions{} + rootCmd = &cobra.Command{ Use: "pc", Short: "Manage your Pinecone vector database infrastructure from the command line", PersistentPreRun: func(cmd *cobra.Command, args []string) { pcio.SetQuiet(globalOptions.quiet) + + // Apply timeout to the command context + if globalOptions.timeout > 0 { + ctx, cancel := context.WithTimeout(cmd.Context(), globalOptions.timeout) + cancelRootFunc = cancel + cmd.SetContext(ctx) + } + }, + PersistentPostRun: func(cmd *cobra.Command, args []string) { + // Cancel the root context when the command completes + if cancelRootFunc != nil { + cancelRootFunc() + cancelRootFunc = nil + } }, Example: help.Examples(` pc login @@ -77,6 +108,7 @@ func init() { Long: rootHelp, } + // Help template and rendering rootCmd.SetHelpTemplate(help.HelpTemplate) help.EnableHelpRendering(rootCmd) @@ -111,4 +143,5 @@ func init() { // Global flags rootCmd.PersistentFlags().BoolVarP(&globalOptions.quiet, "quiet", "q", false, "suppress output") + rootCmd.PersistentFlags().DurationVar(&globalOptions.timeout, "timeout", defaultTimeout, "timeout for commands, defaults to 60s (0 to disable)") } diff --git a/internal/pkg/cli/command/target/target.go b/internal/pkg/cli/command/target/target.go index 17c002d..c562a48 100644 --- a/internal/pkg/cli/command/target/target.go +++ b/internal/pkg/cli/command/target/target.go @@ -71,6 +71,7 @@ func NewTargetCmd() *cobra.Command { Example: targetExample, GroupID: help.GROUP_AUTH.ID, Run: func(cmd *cobra.Command, args []string) { + ctx := cmd.Context() log.Debug(). Str("org", options.org). Str("project", options.project). @@ -131,7 +132,7 @@ func NewTargetCmd() *cobra.Command { exit.ErrorMsg("You must be logged in or have service account credentials configured to set a target context") } - ac := sdk.NewPineconeAdminClient() + ac := sdk.NewPineconeAdminClient(ctx) // Fetch the user's organizations orgs, err := ac.Organization.List(cmd.Context()) @@ -157,7 +158,7 @@ func NewTargetCmd() *cobra.Command { // If the org chosen differs from the current orgId in the token, we need to login again if currentTokenOrgId != "" && currentTokenOrgId != targetOrg.Id { oauth.Logout() - err = login.GetAndSetAccessToken(&targetOrg.Id) + err = login.GetAndSetAccessToken(ctx, &targetOrg.Id) if err != nil { msg.FailMsg("Failed to get access token: %s", err) exit.Error(err, "Error getting access token") @@ -165,7 +166,7 @@ func NewTargetCmd() *cobra.Command { } } - ac := sdk.NewPineconeAdminClient() + ac := sdk.NewPineconeAdminClient(ctx) // Fetch the user's projects projects, err := ac.Project.List(cmd.Context()) if err != nil { @@ -203,7 +204,7 @@ func NewTargetCmd() *cobra.Command { // If the org chosen differs from the current orgId in the token, we need to login again if currentTokenOrgId != org.Id { oauth.Logout() - err = login.GetAndSetAccessToken(&org.Id) + err = login.GetAndSetAccessToken(ctx, &org.Id) if err != nil { msg.FailMsg("Failed to get access token: %s", err) exit.Error(err, "Error getting access token") @@ -229,7 +230,7 @@ func NewTargetCmd() *cobra.Command { if options.project != "" || options.projectID != "" { // We need to reinstantiate the admin client to ensure any auth changes that have happened above // are properly reflected - ac := sdk.NewPineconeAdminClient() + ac := sdk.NewPineconeAdminClient(ctx) // Fetch the user's projects projects, err := ac.Project.List(cmd.Context()) diff --git a/internal/pkg/utils/bodyutil/body.go b/internal/pkg/utils/bodyutil/body.go new file mode 100644 index 0000000..058ca84 --- /dev/null +++ b/internal/pkg/utils/bodyutil/body.go @@ -0,0 +1,55 @@ +package bodyutil + +import ( + "encoding/json" + "fmt" + "os" + + "github.com/pinecone-io/cli/internal/pkg/utils/stdin" +) + +type ArgSource int + +const ( + SourceInline ArgSource = iota + SourceFile + SourceStdin +) + +type SourceInfo struct { + Kind ArgSource `json:"kind"` // SourceInline, SourceFile, SourceStdin + Label string `json:"label"` // "inline", file path, "stdin" +} + +// DecodeBodyArgs unmarshals a JSON body argument (inline/@file/@-) into the provided generic type. +func DecodeBodyArgs[T any](spec string) (*T, SourceInfo, error) { + b, src, err := ReadArg(spec) + if err != nil || len(b) == 0 { + return nil, src, err + } + var out T + if err := json.Unmarshal(b, &out); err != nil { + return nil, src, fmt.Errorf("invalid JSON from %s: %w", src.Label, err) + } + return &out, src, nil +} + +// ReadArg reads inline JSON, @file, or @- (stdin once) and returns bytes and a source label. +func ReadArg(spec string) ([]byte, SourceInfo, error) { + switch { + case spec == "": + return nil, SourceInfo{Kind: SourceInline, Label: "inline"}, nil + case spec == "@-": + b, err := stdin.ReadAllOnce() + if err != nil { + return nil, SourceInfo{Kind: SourceStdin, Label: "stdin"}, fmt.Errorf("stdin already consumed; only one argument may use stdin") + } + return b, SourceInfo{Kind: SourceStdin, Label: "stdin"}, nil + case len(spec) > 0 && spec[0] == '@': + path := spec[1:] + b, err := os.ReadFile(path) + return b, SourceInfo{Kind: SourceFile, Label: path}, err + default: + return []byte(spec), SourceInfo{Kind: SourceInline, Label: "inline"}, nil + } +} diff --git a/internal/pkg/utils/flags/flags.go b/internal/pkg/utils/flags/flags.go new file mode 100644 index 0000000..08a0c57 --- /dev/null +++ b/internal/pkg/utils/flags/flags.go @@ -0,0 +1,213 @@ +package flags + +import ( + "encoding/json" + "io" + "maps" + "os" + "strings" + + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/stdin" +) + +type JSONObject map[string]any +type Float32List []float32 +type UInt32List []uint32 +type StringList []string + +func (m *JSONObject) Set(value string) error { + // allow passing "@file.json" to read a file path and parse as JSON + if strings.HasPrefix(value, "@") { + filePath := strings.TrimPrefix(value, "@") + var raw []byte + var err error + if filePath == "-" { + raw, err = stdin.ReadAllOnce() + if err != nil { + if err == io.ErrUnexpectedEOF { + return pcio.Errorf("stdin already consumed; only one argument may use stdin") + } + return err + } + } else { + raw, err = os.ReadFile(filePath) + if err != nil { + return err + } + } + return json.Unmarshal(raw, m) + } + + var tmp map[string]any + if err := json.Unmarshal([]byte(value), &tmp); err != nil { + return pcio.Errorf("failed to parse JSON: %w", err) + } + if *m == nil { + *m = make(map[string]any) + } + maps.Copy((*m), tmp) + return nil +} + +func (m *JSONObject) String() string { + if m == nil || len(*m) == 0 { + return "" + } + b, _ := json.Marshal(m) + return string(b) +} + +func (*JSONObject) Type() string { return "json-object" } + +func (m *Float32List) Set(value string) error { + // allow passing "@file.json" to read a file path and parse as JSON + if strings.HasPrefix(value, "@") { + filePath := strings.TrimPrefix(value, "@") + var raw []byte + var err error + if filePath == "-" { + raw, err = stdin.ReadAllOnce() + if err != nil { + if err == io.ErrUnexpectedEOF { + return pcio.Errorf("stdin already consumed; only one argument may use stdin") + } + return err + } + } else { + raw, err = os.ReadFile(filePath) + if err != nil { + return err + } + } + value = string(raw) + } + + value = strings.TrimSpace(value) + if value == "" { + *m = (*m)[:0] + return nil + } + + // Require JSON array (no CSV/whitespace parsing) + if !strings.HasPrefix(value, "[") || !strings.HasSuffix(value, "]") { + return pcio.Errorf("expected JSON array for float32 list") + } + + var arr []float32 + if err := json.Unmarshal([]byte(value), &arr); err != nil { + return pcio.Errorf("failed to parse JSON float32 array: %w", err) + } + *m = append((*m)[:0], arr...) + return nil +} + +func (m *Float32List) String() string { + if m == nil || len(*m) == 0 { + return "" + } + b, _ := json.Marshal(m) + return string(b) +} + +func (*Float32List) Type() string { return "float32 json-array" } + +func (m *UInt32List) Set(value string) error { + // allow passing "@file.json" to read a file path and parse as JSON + if strings.HasPrefix(value, "@") { + filePath := strings.TrimPrefix(value, "@") + var raw []byte + var err error + if filePath == "-" { + raw, err = stdin.ReadAllOnce() + if err != nil { + if err == io.ErrUnexpectedEOF { + return pcio.Errorf("stdin already consumed; only one argument may use stdin") + } + return err + } + } else { + raw, err = os.ReadFile(filePath) + if err != nil { + return err + } + } + return json.Unmarshal(raw, m) + } + + value = strings.TrimSpace(value) + if value == "" { + *m = (*m)[:0] + return nil + } + + // Require JSON array (no CSV/whitespace parsing) + if !strings.HasPrefix(value, "[") || !strings.HasSuffix(value, "]") { + return pcio.Errorf("expected JSON array for uint32 list") + } + + var arr []uint32 + if err := json.Unmarshal([]byte(value), &arr); err != nil { + return pcio.Errorf("failed to parse JSON uint32 array: %w", err) + } + *m = append((*m)[:0], arr...) + return nil +} + +func (m *UInt32List) String() string { + if m == nil || len(*m) == 0 { + return "" + } + b, _ := json.Marshal(m) + return string(b) +} + +func (*UInt32List) Type() string { return "uint32 json-array" } + +func (m *StringList) Set(value string) error { + // allow passing "@file.json" to read a file path and parse as JSON + if strings.HasPrefix(value, "@") { + filePath := strings.TrimPrefix(value, "@") + var raw []byte + var err error + if filePath == "-" { + raw, err = stdin.ReadAllOnce() + if err != nil { + if err == io.ErrUnexpectedEOF { + return pcio.Errorf("stdin already consumed; only one argument may use stdin") + } + return err + } + } else { + raw, err = os.ReadFile(filePath) + if err != nil { + return err + } + } + value = string(raw) + } + value = strings.TrimSpace(value) + if value == "" { + *m = (*m)[:0] + return nil + } + if !strings.HasPrefix(value, "[") || !strings.HasSuffix(value, "]") { + return pcio.Errorf("expected JSON array for string list") + } + var arr []string + if err := json.Unmarshal([]byte(value), &arr); err != nil { + return pcio.Errorf("failed to parse JSON string array: %w", err) + } + *m = append((*m)[:0], arr...) + return nil +} + +func (m *StringList) String() string { + if m == nil || len(*m) == 0 { + return "" + } + b, _ := json.Marshal(m) + return string(b) +} + +func (*StringList) Type() string { return "string json-array" } diff --git a/internal/pkg/utils/login/login.go b/internal/pkg/utils/login/login.go index 91ff2ad..4d030a1 100644 --- a/internal/pkg/utils/login/login.go +++ b/internal/pkg/utils/login/login.go @@ -43,7 +43,7 @@ type IO struct { type Options struct{} func Run(ctx context.Context, io IO, opts Options) { - err := GetAndSetAccessToken(nil) + err := GetAndSetAccessToken(ctx, nil) if err != nil { msg.FailMsg("Error acquiring access token while logging in: %s", err) exit.Error(err, "Error acquiring access token while logging in") @@ -62,7 +62,7 @@ func Run(ctx context.Context, io IO, opts Options) { } msg.SuccessMsg("Logged in as " + style.Emphasis(claims.Email) + ". Defaulted to organization ID: " + style.Emphasis(claims.OrgId)) - ac := sdk.NewPineconeAdminClient() + ac := sdk.NewPineconeAdminClient(ctx) if err != nil { msg.FailMsg("Error creating Pinecone admin client: %s", err) exit.Error(err, "Error creating Pinecone admin client") @@ -121,8 +121,7 @@ func Run(ctx context.Context, io IO, opts Options) { // Takes an optional orgId, and attempts to acquire an access token scoped to the orgId if provided. // If a token is successfully acquired it's set in the secrets store, and the user is considered logged in with state.AuthUserToken. -func GetAndSetAccessToken(orgId *string) error { - ctx := context.Background() +func GetAndSetAccessToken(ctx context.Context, orgId *string) error { a := oauth.Auth{} // CSRF state diff --git a/internal/pkg/utils/presenters/describe_index_stats.go b/internal/pkg/utils/presenters/describe_index_stats.go new file mode 100644 index 0000000..422a1ba --- /dev/null +++ b/internal/pkg/utils/presenters/describe_index_stats.go @@ -0,0 +1,42 @@ +package presenters + +import ( + "strings" + + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +func PrintDescribeIndexStatsTable(resp *pinecone.DescribeIndexStatsResponse) { + writer := NewTabWriter() + + columns := []string{"ATTRIBUTE", "VALUE"} + header := strings.Join(columns, "\t") + "\n" + pcio.Fprint(writer, header) + + dimension := uint32(0) + if resp.Dimension != nil { + dimension = *resp.Dimension + } + + pcio.Fprintf(writer, "Dimension\t%d\n", dimension) + pcio.Fprintf(writer, "Index Fullness\t%f\n", resp.IndexFullness) + pcio.Fprintf(writer, "Total Vector Count\t%d\n", resp.TotalVectorCount) + + if len(resp.Namespaces) == 0 { + pcio.Fprintf(writer, "Namespaces\t\n") + } else { + pcio.Fprintf(writer, "Namespaces\t\n") + pcio.Fprintf(writer, "\tNAME\tVECTOR COUNT\n") + + names := make([]string, 0, len(resp.Namespaces)) + for name := range resp.Namespaces { + names = append(names, name) + } + for _, name := range names { + pcio.Fprintf(writer, "\t%s\t%d\n", name, resp.Namespaces[name].VectorCount) + } + } + + writer.Flush() +} diff --git a/internal/pkg/utils/presenters/fetch_vectors.go b/internal/pkg/utils/presenters/fetch_vectors.go new file mode 100644 index 0000000..84b556f --- /dev/null +++ b/internal/pkg/utils/presenters/fetch_vectors.go @@ -0,0 +1,117 @@ +package presenters + +import ( + "sort" + "strings" + + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +type FetchVectorsResults struct { + Vectors map[string]*pinecone.Vector `json:"vectors,omitempty"` + Namespace string `json:"namespace"` + Usage *pinecone.Usage `json:"usage,omitempty"` + Pagination *pinecone.Pagination `json:"pagination,omitempty"` +} + +func NewFetchVectorsResultsFromFetch(resp *pinecone.FetchVectorsResponse) *FetchVectorsResults { + if resp == nil { + return &FetchVectorsResults{} + } + return &FetchVectorsResults{ + Vectors: resp.Vectors, + Namespace: resp.Namespace, + Usage: resp.Usage, + } +} + +func NewFetchVectorsResultsFromFetchByMetadata(resp *pinecone.FetchVectorsByMetadataResponse) *FetchVectorsResults { + if resp == nil { + return &FetchVectorsResults{} + } + return &FetchVectorsResults{ + Vectors: resp.Vectors, + Namespace: resp.Namespace, + Usage: resp.Usage, + Pagination: resp.Pagination, + } +} + +func PrintFetchVectorsTable(results *FetchVectorsResults) { + writer := NewTabWriter() + + // Header Block + if results.Namespace != "" { + pcio.Fprintf(writer, "Namespace: %s\n", results.Namespace) + } + if results.Usage != nil { + pcio.Fprintf(writer, "Usage: %d (read units)\n", results.Usage.ReadUnits) + } + + // Table Header + columns := []string{"ID", "DIMENSION", "VALUES", "SPARSE VALUES", "METADATA"} + pcio.Fprintln(writer, strings.Join(columns, "\t")) + + // Rows + for id, vector := range results.Vectors { + dim := 0 + if vector.Values != nil { + dim = len(*vector.Values) + } + sparseDim := 0 + if vector.SparseValues != nil { + sparseDim = len(vector.SparseValues.Values) + } + metadata := "" + if vector.Metadata != nil { + m := vector.Metadata.AsMap() + if len(m) > 0 { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + show := keys + if len(show) > 3 { + show = show[:3] + } + limited := make(map[string]any, len(show)) + for _, k := range show { + limited[k] = m[k] + } + + s := text.InlineJSON(limited) // compact one-line JSON + if len(keys) > 3 { + // put ellipsis inside the braces: {"a":1,"b":2,"c":3, ...} + s = strings.TrimRight(s, "}") + ", ...}" + } + metadata = s + } + } + + preview := previewSliceFloat32(vector.Values, 3) + row := []string{id, pcio.Sprintf("%d", dim), preview, pcio.Sprintf("%d", sparseDim), metadata} + pcio.Fprintln(writer, strings.Join(row, "\t")) + } + + writer.Flush() +} + +func previewSliceFloat32(values *[]float32, limit int) string { + if values == nil || len(*values) == 0 { + return "" + } + vals := *values + truncated := false + if len(vals) > limit { + vals = vals[:limit] + truncated = true + } + text := text.InlineJSON(vals) + if truncated && strings.HasSuffix(text, "]") { + text = text[:len(text)-1] + ", ...]" + } + return text +} diff --git a/internal/pkg/utils/presenters/list_vectors.go b/internal/pkg/utils/presenters/list_vectors.go new file mode 100644 index 0000000..2c403e0 --- /dev/null +++ b/internal/pkg/utils/presenters/list_vectors.go @@ -0,0 +1,40 @@ +package presenters + +import ( + "strings" + + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +func PrintListVectorsTable(resp *pinecone.ListVectorsResponse) { + writer := NewTabWriter() + + // Header block + if resp.Namespace != "" { + pcio.Fprintf(writer, "Namespace: %s\n", resp.Namespace) + } + if resp.Usage != nil { + pcio.Fprintf(writer, "Usage: %d (read units)\n", resp.Usage.ReadUnits) + } + + // Table header + columns := []string{"ID"} + pcio.Fprintln(writer, strings.Join(columns, "\t")) + + // Rows + for _, vectorId := range resp.VectorIds { + id := "" + if vectorId != nil { + id = *vectorId + } + pcio.Fprintln(writer, id) + } + + // Pagination footer + if resp.NextPaginationToken != nil { + pcio.Fprintf(writer, "Next pagination token: %s\n", *resp.NextPaginationToken) + } + + writer.Flush() +} diff --git a/internal/pkg/utils/presenters/managed_keys.go b/internal/pkg/utils/presenters/managed_keys.go deleted file mode 100644 index 0cc9e34..0000000 --- a/internal/pkg/utils/presenters/managed_keys.go +++ /dev/null @@ -1 +0,0 @@ -package presenters diff --git a/internal/pkg/utils/presenters/query_vectors.go b/internal/pkg/utils/presenters/query_vectors.go new file mode 100644 index 0000000..271101e --- /dev/null +++ b/internal/pkg/utils/presenters/query_vectors.go @@ -0,0 +1,131 @@ +package presenters + +import ( + "sort" + "strings" + + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/cli/internal/pkg/utils/text" + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +func PrintQueryVectorsTable(resp *pinecone.QueryVectorsResponse) { + writer := NewTabWriter() + + // Header Block + if resp.Namespace != "" { + pcio.Fprintf(writer, "Namespace: %s\n", resp.Namespace) + } + if resp.Usage != nil { + pcio.Fprintf(writer, "Usage: %d (read units)\n", resp.Usage.ReadUnits) + } + + // Detect which columns to show + hasDense := false + hasSparse := false + hasMetadata := false + for _, m := range resp.Matches { + if m == nil || m.Vector == nil { + continue + } + if m.Vector.Values != nil && len(*m.Vector.Values) > 0 { + hasDense = true + } + if m.Vector.SparseValues != nil && + (len(m.Vector.SparseValues.Indices) > 0 || len(m.Vector.SparseValues.Values) > 0) { + hasSparse = true + } + if m.Vector.Metadata != nil { + hasMetadata = true + } + } + + // Table Header + cols := []string{"ID", "SCORE"} + if hasDense { + cols = append(cols, "VALUES") + } + if hasSparse { + cols = append(cols, "SPARSE INDICES", "SPARSE VALUES") + } + if hasMetadata { + cols = append(cols, "METADATA") + } + pcio.Fprintln(writer, strings.Join(cols, "\t")) + + // Rows + for _, match := range resp.Matches { + if match == nil || match.Vector == nil { + continue + } + row := []string{match.Vector.Id, pcio.Sprintf("%f", match.Score)} + + if hasDense { + values := "" + if match.Vector.Values != nil { + values = previewSliceFloat32(match.Vector.Values, 3) + } + row = append(row, values) + } + if hasSparse { + iPreview, vPreview := "", "" + if match.Vector.SparseValues != nil { + iPreview = previewSliceUint32(match.Vector.SparseValues.Indices, 3) + vPreview = previewSliceFloat32(&match.Vector.SparseValues.Values, 3) + } + row = append(row, iPreview, vPreview) + } + if hasMetadata { + metadata := "" + if match.Vector.Metadata != nil { + m := match.Vector.Metadata.AsMap() + if len(m) > 0 { + keys := make([]string, 0, len(m)) + for k := range m { + keys = append(keys, k) + } + sort.Strings(keys) + show := keys + if len(show) > 3 { + show = show[:3] + } + limited := make(map[string]any, len(show)) + for _, k := range show { + limited[k] = m[k] + } + + s := text.InlineJSON(limited) // compact one-line JSON + if len(keys) > 3 { + // put ellipsis inside the braces: {"a":1,"b":2,"c":3, ...} + s = strings.TrimRight(s, "}") + ", ...}" + } + metadata = s + } + } + row = append(row, metadata) + } + + pcio.Fprintln(writer, strings.Join(row, "\t")) + } + + writer.Flush() +} + +func previewSliceUint32(values []uint32, limit int) string { + if len(values) == 0 { + return "" + } + vals := values + truncated := false + if len(vals) > limit { + vals = vals[:limit] + truncated = true + } + + text := text.InlineJSON(vals) + if truncated && strings.HasSuffix(text, "]") { + text = text[:len(text)-1] + ", ...]" + } + + return text +} diff --git a/internal/pkg/utils/presenters/update_vector.go b/internal/pkg/utils/presenters/update_vector.go new file mode 100644 index 0000000..5c9ede9 --- /dev/null +++ b/internal/pkg/utils/presenters/update_vector.go @@ -0,0 +1,20 @@ +package presenters + +import ( + "strings" + + "github.com/pinecone-io/cli/internal/pkg/utils/pcio" + "github.com/pinecone-io/go-pinecone/v5/pinecone" +) + +func PrintUpdateVectorsByMetadataTable(resp *pinecone.UpdateVectorsByMetadataResponse) { + writer := NewTabWriter() + + columns := []string{"ATTRIBUTE", "VALUE"} + header := strings.Join(columns, "\t") + "\n" + pcio.Fprint(writer, header) + + pcio.Fprintf(writer, "Matched Records\t%d\n", resp.MatchedRecords) + + writer.Flush() +} diff --git a/internal/pkg/utils/sdk/client.go b/internal/pkg/utils/sdk/client.go index 1acf11e..daf916b 100644 --- a/internal/pkg/utils/sdk/client.go +++ b/internal/pkg/utils/sdk/client.go @@ -24,7 +24,7 @@ const ( CLISourceTag = "pinecone-cli" ) -func NewPineconeClient() *pinecone.Client { +func NewPineconeClient(ctx context.Context) *pinecone.Client { targetOrgId := state.TargetOrg.Get().Id targetProjectId := state.TargetProj.Get().Id log.Debug(). @@ -32,7 +32,6 @@ func NewPineconeClient() *pinecone.Client { Str("targetProjectId", targetProjectId). Msg("Loading target context") - ctx := context.Background() oauth2Token, err := oauth.Token(ctx) if err != nil { log.Error().Err(err).Msg("Error retrieving oauth token") @@ -65,12 +64,11 @@ func NewPineconeClient() *pinecone.Client { exit.ErrorMsg("No target project set") } - return NewPineconeClientForProjectById(targetProjectId) + return NewPineconeClientForProjectById(ctx, targetProjectId) } -func NewPineconeClientForProjectById(projectId string) *pinecone.Client { - ac := NewPineconeAdminClient() - ctx := context.Background() +func NewPineconeClientForProjectById(ctx context.Context, projectId string) *pinecone.Client { + ac := NewPineconeAdminClient(ctx) project, err := ac.Project.Describe(ctx, projectId) if err != nil { @@ -121,8 +119,7 @@ func NewClientForAPIKey(apiKey string) *pinecone.Client { return pc } -func NewPineconeAdminClient() *pinecone.AdminClient { - ctx := context.Background() +func NewPineconeAdminClient(ctx context.Context) *pinecone.AdminClient { oauth2Token, err := oauth.Token(ctx) if err != nil { log.Error().Err(err).Msg("Error retrieving oauth token") @@ -153,6 +150,22 @@ func NewPineconeAdminClient() *pinecone.AdminClient { return ac } +func NewIndexConnection(ctx context.Context, pc *pinecone.Client, indexName string, namespace string) (*pinecone.IndexConnection, error) { + index, err := pc.DescribeIndex(ctx, indexName) + if err != nil { + return nil, pcio.Errorf("failed to describe index: %w", err) + } + + ic, err := pc.Index(pinecone.NewIndexConnParams{ + Host: index.Host, + Namespace: namespace, + }) + if err != nil { + return nil, pcio.Errorf("failed to create index connection: %w", err) + } + return ic, nil +} + func getCLIAPIKeyForProject(ctx context.Context, ac *pinecone.AdminClient, project *pinecone.Project) (string, error) { projectAPIKeysMap := secrets.ManagedAPIKeys.Get() var managedKey secrets.ManagedKey diff --git a/internal/pkg/utils/stdin/reader.go b/internal/pkg/utils/stdin/reader.go new file mode 100644 index 0000000..01f94dd --- /dev/null +++ b/internal/pkg/utils/stdin/reader.go @@ -0,0 +1,24 @@ +package stdin + +import ( + "io" + "os" + "sync/atomic" +) + +var consumed atomic.Bool + +// ReadAllOnce reads all data from stdin exactly once across the process. +// If stdin was already consumed by another reader, it returns an error. +func ReadAllOnce() ([]byte, error) { + if !consumed.CompareAndSwap(false, true) { + return nil, io.ErrUnexpectedEOF + } + return io.ReadAll(os.Stdin) +} + +// HasPipedStdin returns true if stdin is a pipe (not a TTY). +func HasPipedStdin() bool { + fi, _ := os.Stdin.Stat() + return fi != nil && (fi.Mode()&os.ModeCharDevice) == 0 +}