Skip to content

Commit 838d864

Browse files
AmperAmper
committed
implemented ability to set API key per-request (necessary in MCP server)
1 parent 3d15656 commit 838d864

File tree

5 files changed

+159
-21
lines changed

5 files changed

+159
-21
lines changed

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
module github.com/VictoriaMetrics/victoriametrics-cloud-api-go
22

3-
go 1.24
3+
go 1.26

‎v1/client.go‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ const (
1414
DefaultBaseURL = "https://api.victoriametrics.cloud"
1515
// AccessTokenHeader is the header name for the access token
1616
AccessTokenHeader = "X-VM-Cloud-Access"
17+
// DynamicAPIKey - use this constant as a value for API key in the context to indicate that the API key
18+
// should be taken from the context instead of the client configuration.
19+
// This allows using different API keys for different requests with the same client instance.
20+
DynamicAPIKey = "dynamic"
1721
)
1822

1923
// VMCloudAPIClient represents a API client for VictoriaMetrics Cloud API
@@ -46,6 +50,9 @@ func New(apiKey string, options ...VMCloudAPIClientOption) (*VMCloudAPIClient, e
4650
if apiKey == "" {
4751
return nil, fmt.Errorf("API key cannot be empty")
4852
}
53+
if apiKey == DynamicAPIKey {
54+
apiKey = ""
55+
}
4956
result := &VMCloudAPIClient{
5057
c: http.DefaultClient,
5158
apiKey: apiKey,

‎v1/client_test.go‎

Lines changed: 28 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,18 @@ import (
1010

1111
func TestNew(t *testing.T) {
1212
tests := []struct {
13-
name string
14-
apiKey string
15-
options []VMCloudAPIClientOption
16-
wantErr bool
13+
name string
14+
apiKey string
15+
wantAPIKey string
16+
options []VMCloudAPIClientOption
17+
wantErr bool
1718
}{
1819
{
19-
name: "valid client with API key",
20-
apiKey: "test-api-key",
21-
options: nil,
22-
wantErr: false,
20+
name: "valid client with API key",
21+
apiKey: "test-api-key",
22+
wantAPIKey: "test-api-key",
23+
options: nil,
24+
wantErr: false,
2325
},
2426
{
2527
name: "empty API key",
@@ -28,8 +30,9 @@ func TestNew(t *testing.T) {
2830
wantErr: true,
2931
},
3032
{
31-
name: "valid client with custom HTTP client",
32-
apiKey: "test-api-key",
33+
name: "valid client with custom HTTP client",
34+
apiKey: "test-api-key",
35+
wantAPIKey: "test-api-key",
3336
options: []VMCloudAPIClientOption{
3437
WithHTTPClient(&http.Client{
3538
Timeout: 30 * time.Second,
@@ -38,21 +41,30 @@ func TestNew(t *testing.T) {
3841
wantErr: false,
3942
},
4043
{
41-
name: "valid client with custom base URL",
42-
apiKey: "test-api-key",
44+
name: "valid client with custom base URL",
45+
apiKey: "test-api-key",
46+
wantAPIKey: "test-api-key",
4347
options: []VMCloudAPIClientOption{
4448
WithBaseURL("https://custom-api.victoriametrics.com"),
4549
},
4650
wantErr: false,
4751
},
4852
{
49-
name: "invalid base URL",
50-
apiKey: "test-api-key",
53+
name: "invalid base URL",
54+
apiKey: "test-api-key",
55+
wantAPIKey: "test-api-key",
5156
options: []VMCloudAPIClientOption{
5257
WithBaseURL("://invalid-url"),
5358
},
5459
wantErr: true,
5560
},
61+
{
62+
name: "dynamic API key stores empty string",
63+
apiKey: DynamicAPIKey,
64+
wantAPIKey: "",
65+
options: nil,
66+
wantErr: false,
67+
},
5668
}
5769

5870
for _, tt := range tests {
@@ -66,8 +78,8 @@ func TestNew(t *testing.T) {
6678
if client == nil {
6779
t.Errorf("New() returned nil client without error")
6880
} else {
69-
if client.apiKey != tt.apiKey {
70-
t.Errorf("New() client.apiKey = %v, want %v", client.apiKey, tt.apiKey)
81+
if client.apiKey != tt.wantAPIKey {
82+
t.Errorf("New() client.apiKey = %v, want %v", client.apiKey, tt.wantAPIKey)
7183
}
7284
}
7385
}

‎v1/utils.go‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,25 +8,39 @@ import (
88
"net/http"
99
)
1010

11+
type apiKeyContextKeyType string
12+
13+
const apiKeyContextKey apiKeyContextKeyType = "X-VM-Cloud-Access"
14+
15+
func ContextWithDynamicAPIKey(ctx context.Context, apiKey string) context.Context {
16+
return context.WithValue(ctx, apiKeyContextKey, apiKey)
17+
}
18+
1119
func requestAPI[R any](ctx context.Context, a *VMCloudAPIClient, method string, body io.Reader, path ...string) (R, error) {
1220
var result R
1321
reqURL := a.parsedURL.JoinPath(path...).String()
1422
req, err := http.NewRequestWithContext(ctx, method, reqURL, body)
1523
if err != nil {
1624
return result, fmt.Errorf("failed to create request: %w", err)
1725
}
18-
req.Header.Set(AccessTokenHeader, a.apiKey)
26+
apiKey := a.apiKey
27+
if apiKey == "" {
28+
if apiKeyFromCtx, ok := ctx.Value(apiKeyContextKey).(string); ok && apiKeyFromCtx != "" {
29+
apiKey = apiKeyFromCtx
30+
}
31+
}
32+
req.Header.Set(AccessTokenHeader, apiKey)
1933
resp, err := a.c.Do(req)
2034
if err != nil {
2135
return result, fmt.Errorf("failed to send request: %w", err)
2236
}
37+
defer func() {
38+
_ = resp.Body.Close()
39+
}()
2340
respBodyBytes, err := io.ReadAll(resp.Body)
2441
if err != nil {
2542
return result, fmt.Errorf("failed to read response body: %w", err)
2643
}
27-
defer func() {
28-
_ = resp.Body.Close()
29-
}()
3044
if resp.StatusCode/100 != 2 {
3145
return result, fmt.Errorf("unexpected status code: %d, body: %s", resp.StatusCode, string(respBodyBytes))
3246
}

‎v1/utils_test.go‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
11
package v1
22

33
import (
4+
"context"
5+
"encoding/json"
6+
"net/http"
7+
"net/http/httptest"
48
"testing"
59
)
610

@@ -74,6 +78,107 @@ func TestCheckDeploymentID(t *testing.T) {
7478
}
7579
}
7680

81+
func TestContextWithDynamicAPIKey(t *testing.T) {
82+
tests := []struct {
83+
name string
84+
apiKey string
85+
}{
86+
{
87+
name: "non-empty key",
88+
apiKey: "my-secret-key",
89+
},
90+
{
91+
name: "empty key",
92+
apiKey: "",
93+
},
94+
}
95+
96+
for _, tt := range tests {
97+
t.Run(tt.name, func(t *testing.T) {
98+
ctx := ContextWithDynamicAPIKey(context.Background(), tt.apiKey)
99+
got, ok := ctx.Value(apiKeyContextKey).(string)
100+
if !ok {
101+
t.Fatal("ContextWithDynamicAPIKey() did not store a string value in context")
102+
}
103+
if got != tt.apiKey {
104+
t.Errorf("ContextWithDynamicAPIKey() stored %q, want %q", got, tt.apiKey)
105+
}
106+
})
107+
}
108+
}
109+
110+
func TestRequestAPIDynamicAPIKey(t *testing.T) {
111+
tests := []struct {
112+
name string
113+
clientKey string
114+
ctxKey *string // nil means don't set context key
115+
expectedHeader string
116+
}{
117+
{
118+
name: "static key, no ctx",
119+
clientKey: "static-key",
120+
ctxKey: nil,
121+
expectedHeader: "static-key",
122+
},
123+
{
124+
name: "static key + ctx key, client takes precedence",
125+
clientKey: "static-key",
126+
ctxKey: strPtr("dynamic-key"),
127+
expectedHeader: "static-key",
128+
},
129+
{
130+
name: "dynamic + ctx key",
131+
clientKey: DynamicAPIKey,
132+
ctxKey: strPtr("per-request-key"),
133+
expectedHeader: "per-request-key",
134+
},
135+
{
136+
name: "dynamic, no ctx key",
137+
clientKey: DynamicAPIKey,
138+
ctxKey: nil,
139+
expectedHeader: "",
140+
},
141+
{
142+
name: "dynamic + empty ctx key",
143+
clientKey: DynamicAPIKey,
144+
ctxKey: strPtr(""),
145+
expectedHeader: "",
146+
},
147+
}
148+
149+
for _, tt := range tests {
150+
t.Run(tt.name, func(t *testing.T) {
151+
var capturedHeader string
152+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
153+
capturedHeader = r.Header.Get(AccessTokenHeader)
154+
w.WriteHeader(http.StatusOK)
155+
_ = json.NewEncoder(w).Encode([]CloudProviderInfo{})
156+
}))
157+
defer server.Close()
158+
159+
client, err := New(tt.clientKey, WithBaseURL(server.URL))
160+
if err != nil {
161+
t.Fatalf("New() error = %v", err)
162+
}
163+
164+
ctx := context.Background()
165+
if tt.ctxKey != nil {
166+
ctx = ContextWithDynamicAPIKey(ctx, *tt.ctxKey)
167+
}
168+
169+
_, _ = client.ListCloudProviders(ctx)
170+
171+
if capturedHeader != tt.expectedHeader {
172+
t.Errorf("request header %q = %q, want %q", AccessTokenHeader, capturedHeader, tt.expectedHeader)
173+
}
174+
})
175+
}
176+
}
177+
178+
func strPtr(s string) *string {
179+
return &s
180+
}
181+
77182
func TestIsValidTenantID(t *testing.T) {
78183
tests := []struct {
79184
name string

0 commit comments

Comments
 (0)