From adbc16e460ec3fd959c467bda1eaae80fdfa87ca Mon Sep 17 00:00:00 2001 From: haruska <2803+haruska@users.noreply.github.com> Date: Fri, 5 Apr 2024 09:17:56 -0500 Subject: [PATCH] add optional SourceTag to ManagementClient --- pinecone/management_client.go | 35 +++++++++++++++--- pinecone/management_client_test.go | 59 ++++++++++++++++++++++++++++-- 2 files changed, 85 insertions(+), 9 deletions(-) diff --git a/pinecone/management_client.go b/pinecone/management_client.go index 83b9465..4cd82fa 100644 --- a/pinecone/management_client.go +++ b/pinecone/management_client.go @@ -6,6 +6,8 @@ import ( "github.com/deepmap/oapi-codegen/v2/pkg/securityprovider" "github.com/google/uuid" "github.com/pinecone-io/go-pinecone/internal/gen/management" + "github.com/pinecone-io/go-pinecone/internal/provider" + "github.com/pinecone-io/go-pinecone/internal/useragent" "net/http" ) @@ -21,10 +23,14 @@ import ( // // Fields: // - apiKey: The API key used for authenticating requests to the management API. -// This key should have the necessary permissions for the operations you intend to perform. +// This key should have the necessary organization-level permissions for the operations +// you intend to perform. // - restClient: An instance of the generated low-level client that actually performs // HTTP requests to the management API. This field is internal and managed // by the ManagementClient. +// - sourceTag: An optional string used to identify the source of the requests. This tag +// is included in the User-Agent header of each request made by this client, providing +// a way to trace and analyze the usage patterns or origins of API requests. // // To use ManagementClient, first instantiate it using the NewManagementClient function, // providing it with the necessary configuration. Once instantiated, you can call its @@ -33,7 +39,10 @@ import ( // // Example: // -// clientParams := NewManagementClientParams{ApiKey: "your_api_key_here"} +// clientParams := NewManagementClientParams{ +// ApiKey: "your_api_key_here", +// SourceTag: "your_source_identifier", // optional +// } // managementClient, err := NewManagementClient(clientParams) // if err != nil { // log.Fatalf("Failed to create management client: %v", err) @@ -43,14 +52,17 @@ import ( // Note that ManagementClient methods are designed to be safe for concurrent use by multiple // goroutines, assuming that its configuration (e.g., the API key) is not modified after // initialization. + type ManagementClient struct { apiKey string restClient *management.ClientWithResponses + sourceTag string } // NewManagementClientParams holds the parameters for creating a new ManagementClient. type NewManagementClientParams struct { - ApiKey string + ApiKey string + SourceTag string // optional } // NewManagementClient creates and initializes a new instance of ManagementClient. @@ -58,7 +70,12 @@ type NewManagementClientParams struct { // authentication and communication with the management API. // // The method requires an input parameter of type NewManagementClientParams, which includes: -// - ApiKey: A string representing the API key used for authentication against the management API. +// - ApiKey: The API key used for authenticating requests to the management API. +// This key should have the necessary organization-level permissions for the operations +// you intend to perform. +// - SourceTag: An optional string used to identify the source of the requests. This tag +// is included in the User-Agent header of each request made by this client, providing +// a way to trace and analyze the usage patterns or origins of API requests. // // The API key is used to configure the underlying HTTP client with the appropriate // authentication headers for all requests made to the management API. @@ -72,6 +89,7 @@ type NewManagementClientParams struct { // // clientParams := NewManagementClientParams{ // ApiKey: "your_api_key_here", +// SourceTag: "my-application", // Optional but recommended for identifying request sources // } // managementClient, err := NewManagementClient(clientParams) // if err != nil { @@ -87,12 +105,17 @@ func NewManagementClient(in NewManagementClientParams) (*ManagementClient, error return nil, err } - client, err := management.NewClientWithResponses("https://api.pinecone.io/management/v1alpha", management.WithRequestEditorFn(apiKeyProvider.Intercept)) + userAgentProvider := provider.NewHeaderProvider("User-Agent", useragent.BuildUserAgent(in.SourceTag)) + + client, err := management.NewClientWithResponses("https://api.pinecone.io/management/v1alpha", + management.WithRequestEditorFn(apiKeyProvider.Intercept), + management.WithRequestEditorFn(userAgentProvider.Intercept), + ) if err != nil { return nil, err } - c := ManagementClient{apiKey: in.ApiKey, restClient: client} + c := ManagementClient{apiKey: in.ApiKey, restClient: client, sourceTag: in.SourceTag} return &c, nil } diff --git a/pinecone/management_client_test.go b/pinecone/management_client_test.go index eb8fa72..942092c 100644 --- a/pinecone/management_client_test.go +++ b/pinecone/management_client_test.go @@ -14,9 +14,11 @@ import ( type ManagementClientTests struct { suite.Suite - client ManagementClient - project Project - apiKey APIKeyWithoutSecret + client ManagementClient + clientSourceTag ManagementClient + project Project + apiKey APIKeyWithoutSecret + sourceTag string } func TestManagementClient(t *testing.T) { @@ -33,6 +35,13 @@ func (ts *ManagementClientTests) SetupSuite() { } ts.client = *client + ts.sourceTag = "test_source_tag" + clientSourceTag, err := NewManagementClient(NewManagementClientParams{ApiKey: apiKey, SourceTag: ts.sourceTag}) + if err != nil { + ts.FailNow(err.Error()) + } + ts.clientSourceTag = *clientSourceTag + projects, err := ts.client.ListProjects(context.Background()) require.NoError(ts.T(), err, "Failed to list projects for test setup") require.Greater(ts.T(), len(projects), 0, "Projects list should not be empty") @@ -45,12 +54,47 @@ func (ts *ManagementClientTests) SetupSuite() { ts.apiKey = *apiKeys[0] } +func (ts *ManagementClientTests) TestNewClientParamsSet() { + apiKey := "test-api-key" + client, err := NewManagementClient(NewManagementClientParams{ApiKey: apiKey}) + if err != nil { + ts.FailNow(err.Error()) + } + if client.apiKey != apiKey { + ts.FailNow(fmt.Sprintf("Expected client to have apiKey '%s', but got '%s'", apiKey, client.apiKey)) + } + if client.sourceTag != "" { + ts.FailNow(fmt.Sprintf("Expected client to have empty sourceTag, but got '%s'", client.sourceTag)) + } +} + +func (ts *ManagementClientTests) TestNewClientParamsSetSourceTag() { + apiKey := "test-api-key" + sourceTag := "test-source-tag" + client, err := NewManagementClient(NewManagementClientParams{ApiKey: apiKey, SourceTag: sourceTag}) + if err != nil { + ts.FailNow(err.Error()) + } + if client.apiKey != apiKey { + ts.FailNow(fmt.Sprintf("Expected client to have apiKey '%s', but got '%s'", apiKey, client.apiKey)) + } + if client.sourceTag != sourceTag { + ts.FailNow(fmt.Sprintf("Expected client to have sourceTag '%s', but got '%s'", sourceTag, client.sourceTag)) + } +} + func (ts *ManagementClientTests) TestListProjects() { projects, err := ts.client.ListProjects(context.Background()) require.NoError(ts.T(), err, "Failed to list projects") require.Greater(ts.T(), len(projects), 0, "Projects list should not be empty") } +func (ts *ManagementClientTests) TestListProjectsSourceTag() { + projects, err := ts.clientSourceTag.ListProjects(context.Background()) + require.NoError(ts.T(), err, "Failed to list projects") + require.Greater(ts.T(), len(projects), 0, "Projects list should not be empty") +} + func (ts *ManagementClientTests) TestFetchProject() { testProjectID := ts.project.Id @@ -110,6 +154,15 @@ func (ts *ManagementClientTests) TestFetchApiKey() { require.Equal(ts.T(), apiKeyDetails.ProjectId, ts.apiKey.ProjectId, "API key's Project ID should match") } +func (ts *ManagementClientTests) TestFetchApiKeySourceTag() { + apiKeyDetails, err := ts.clientSourceTag.FetchApiKey(context.Background(), ts.apiKey.Id) + require.NoError(ts.T(), err, "Failed to fetch API key details") + require.NotNil(ts.T(), apiKeyDetails, "API key details should not be nil") + require.Equal(ts.T(), apiKeyDetails.Id, ts.apiKey.Id, "API key ID should match") + require.Equal(ts.T(), apiKeyDetails.Name, ts.apiKey.Name, "API key Name should match") + require.Equal(ts.T(), apiKeyDetails.ProjectId, ts.apiKey.ProjectId, "API key's Project ID should match") +} + func (ts *ManagementClientTests) TestCreateApiKey() { apiKeyName := generateRandomString(6) // current limitation of Alpha-release newApiKey, err := ts.client.CreateApiKey(context.Background(), ts.project.Id, apiKeyName)