Skip to content

Commit 08fb1c9

Browse files
EspenAlbertEspenAlbert
committed
refactor: replace simple context key with structured user agent extra info
1 parent 597d3e2 commit 08fb1c9

File tree

3 files changed

+255
-23
lines changed

3 files changed

+255
-23
lines changed

‎internal/config/transport.go‎

Lines changed: 89 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,112 @@
11
package config
22

33
import (
4+
"context"
45
"fmt"
56
"log"
67
"net/http"
78
"strings"
89
"time"
910
)
1011

11-
type ContextKey string
12+
// UserAgentExtra holds additional metadata to be appended to the User-Agent header and context.
13+
type UserAgentExtra struct {
14+
Type string // Type of the operation (e.g., "Resource", "Datasource", etc.)
15+
Name string // Full name, for example mongodbatlas_database_user
16+
Operation string // GrpcCall for example, ReadResource, see wrapped_provider_server.go for details
17+
ScriptLocation string // TODO: Support setting this field as opt-in on resources and datasources
18+
}
19+
20+
// Combine returns a new UserAgentExtra by merging the receiver with another.
21+
// Non-empty fields in 'other' take precedence over the receiver's fields.
22+
func (e UserAgentExtra) Combine(other UserAgentExtra) UserAgentExtra {
23+
typeName := e.Type
24+
if other.Type != "" {
25+
typeName = other.Type
26+
}
27+
name := e.Name
28+
if other.Name != "" {
29+
name = other.Name
30+
}
31+
operation := e.Operation
32+
if other.Operation != "" {
33+
operation = other.Operation
34+
}
35+
scriptLocation := e.ScriptLocation
36+
if other.ScriptLocation != "" {
37+
scriptLocation = other.ScriptLocation
38+
}
39+
return UserAgentExtra{
40+
Type: typeName,
41+
Name: name,
42+
Operation: operation,
43+
ScriptLocation: scriptLocation,
44+
}
45+
}
46+
47+
// ToHeaderValue returns a string representation suitable for use as a User-Agent header value.
48+
// If oldHeader is non-empty, it is prepended to the new value.
49+
func (e UserAgentExtra) ToHeaderValue(oldHeader string) string {
50+
parts := []string{}
51+
addPart := func(key, part string) {
52+
if part == "" {
53+
return
54+
}
55+
parts = append(parts, fmt.Sprintf("%s/%s", key, part))
56+
}
57+
addPart("Type", e.Type)
58+
addPart("Name", e.Name)
59+
addPart("Operation", e.Operation)
60+
addPart("ScriptLocation", e.ScriptLocation)
61+
newPart := strings.Join(parts, " ")
62+
if oldHeader == "" {
63+
return newPart
64+
}
65+
return fmt.Sprintf("%s %s", oldHeader, newPart)
66+
}
67+
68+
type UserAgentKey string
1269

1370
const (
14-
ContextKeyTFSrc = ContextKey("tf-src")
15-
UserAgentHeader = "User-Agent"
71+
UserAgentExtraKey = UserAgentKey("user-agent-extra")
72+
UserAgentHeader = "User-Agent"
1673
)
1774

75+
// ReadUserAgentExtra retrieves the UserAgentExtra from the context if present.
76+
// Returns a pointer to the UserAgentExtra, or nil if not set or of the wrong type.
77+
// Logs a warning if the value is not of the expected type.
78+
func ReadUserAgentExtra(ctx context.Context) *UserAgentExtra {
79+
extra := ctx.Value(UserAgentExtraKey)
80+
if extra == nil {
81+
return nil
82+
}
83+
if userAgentExtra, ok := extra.(UserAgentExtra); ok {
84+
return &userAgentExtra
85+
}
86+
log.Printf("[WARN] UserAgentExtra in context is not of type UserAgentExtra, got %v", extra)
87+
return nil
88+
}
89+
90+
// AddUserAgentExtra returns a new context with UserAgentExtra merged into any existing value.
91+
// If a UserAgentExtra is already present in the context, the fields of 'extra' will override non-empty fields.
92+
func AddUserAgentExtra(ctx context.Context, extra UserAgentExtra) context.Context {
93+
oldExtra := ReadUserAgentExtra(ctx)
94+
if oldExtra == nil {
95+
return context.WithValue(ctx, UserAgentExtraKey, extra)
96+
}
97+
newExtra := oldExtra.Combine(extra)
98+
return context.WithValue(ctx, UserAgentExtraKey, newExtra)
99+
}
100+
18101
type TFSrcUserAgentAdder struct {
19102
Transport http.RoundTripper
20103
}
21104

22105
func (t *TFSrcUserAgentAdder) RoundTrip(req *http.Request) (*http.Response, error) {
23-
ctx := req.Context()
24-
tfSrcName := ctx.Value(ContextKeyTFSrc)
25-
if tfSrcName != nil {
106+
extra := ReadUserAgentExtra(req.Context())
107+
if extra != nil {
26108
userAgent := req.Header.Get(UserAgentHeader)
27-
tfSrcValue := tfSrcName.(string)
28-
newVar := fmt.Sprintf("%s %s/%s", userAgent, ContextKeyTFSrc, tfSrcValue)
109+
newVar := extra.ToHeaderValue(userAgent)
29110
req.Header.Set(UserAgentHeader, newVar)
30111
}
31112
resp, err := t.Transport.RoundTrip(req)

‎internal/config/transport_test.go‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -151,6 +151,95 @@ func TestNetworkLoggingTransport_Disabled(t *testing.T) {
151151
assert.Empty(t, logStr, "Expected no logs when network logging is disabled")
152152
}
153153

154+
func TestUserAgentExtra_ToHeaderValue(t *testing.T) {
155+
testCases := map[string]struct {
156+
extra config.UserAgentExtra
157+
old string
158+
expected string
159+
}{
160+
"all fields": {
161+
extra: config.UserAgentExtra{
162+
Type: "type1",
163+
Name: "name1",
164+
Operation: "op1",
165+
ScriptLocation: "loc1",
166+
},
167+
old: "base/1.0",
168+
expected: "base/1.0 Type/type1 Name/name1 Operation/op1 ScriptLocation/loc1",
169+
},
170+
"some fields empty": {
171+
extra: config.UserAgentExtra{
172+
Type: "",
173+
Name: "name2",
174+
Operation: "",
175+
},
176+
old: "",
177+
expected: "Name/name2",
178+
},
179+
"none": {
180+
extra: config.UserAgentExtra{},
181+
old: "",
182+
expected: "",
183+
},
184+
}
185+
186+
for name, tc := range testCases {
187+
t.Run(name, func(t *testing.T) {
188+
got := tc.extra.ToHeaderValue(tc.old)
189+
assert.Equal(t, tc.expected, got)
190+
})
191+
}
192+
}
193+
194+
func TestUserAgentExtra_Combine(t *testing.T) {
195+
testCases := map[string]struct {
196+
base config.UserAgentExtra
197+
other config.UserAgentExtra
198+
expected config.UserAgentExtra
199+
}{
200+
"other overwrites non-empty": {
201+
base: config.UserAgentExtra{Type: "A", Name: "B", Operation: "C", ScriptLocation: "D"},
202+
other: config.UserAgentExtra{Type: "X", Name: "Y", Operation: "Z", ScriptLocation: "Q"},
203+
expected: config.UserAgentExtra{Type: "X", Name: "Y", Operation: "Z", ScriptLocation: "Q"},
204+
},
205+
"other empty": {
206+
base: config.UserAgentExtra{Type: "A", Name: "B", Operation: "C", ScriptLocation: "D"},
207+
other: config.UserAgentExtra{},
208+
expected: config.UserAgentExtra{Type: "A", Name: "B", Operation: "C", ScriptLocation: "D"},
209+
},
210+
"mixed": {
211+
base: config.UserAgentExtra{Type: "A", Name: "B"},
212+
other: config.UserAgentExtra{Name: "Y", ScriptLocation: "Q"},
213+
expected: config.UserAgentExtra{Type: "A", Name: "Y", Operation: "", ScriptLocation: "Q"},
214+
},
215+
}
216+
for name, tc := range testCases {
217+
t.Run(name, func(t *testing.T) {
218+
got := tc.base.Combine(tc.other)
219+
assert.Equal(t, tc.expected, got)
220+
})
221+
}
222+
}
223+
224+
func TestAddUserAgentExtra(t *testing.T) {
225+
base := config.UserAgentExtra{Type: "A", Name: "B"}
226+
other := config.UserAgentExtra{Name: "Y", ScriptLocation: "Q"}
227+
ctx := config.AddUserAgentExtra(t.Context(), base)
228+
ctx2 := config.AddUserAgentExtra(ctx, other)
229+
// Should combine base and other
230+
e := ctx2.Value(config.UserAgentExtraKey)
231+
assert.NotNil(t, e)
232+
ua := config.UserAgentExtra{}
233+
if v, ok := e.(config.UserAgentExtra); ok {
234+
ua = v
235+
}
236+
// The combined should have Type from base, Name from other, ScriptLocation from other
237+
assert.Equal(t, "A", ua.Type)
238+
assert.Equal(t, "Y", ua.Name)
239+
assert.Equal(t, "Q", ua.ScriptLocation)
240+
assert.Empty(t, ua.Operation)
241+
}
242+
154243
func TestAccNetworkLogging(t *testing.T) {
155244
acc.SkipInUnitTest(t)
156245
acc.PreCheckBasic(t)

‎internal/provider/wrapper_provider_server.go‎

Lines changed: 77 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ import (
77
"github.com/mongodb/terraform-provider-mongodbatlas/internal/config"
88
)
99

10+
// NewWrappedProviderServer returns a new ProviderServer that wraps the old one.
11+
// This is used to add additional metadata to the User-Agent header and context.
1012
func NewWrappedProviderServer(old func() tfprotov6.ProviderServer) func() tfprotov6.ProviderServer {
1113
return func() tfprotov6.ProviderServer {
1214
return &WrappedProviderServer{
@@ -44,57 +46,101 @@ func (s *WrappedProviderServer) StopProvider(ctx context.Context, req *tfprotov6
4446
}
4547

4648
func (s *WrappedProviderServer) ValidateResourceConfig(ctx context.Context, req *tfprotov6.ValidateResourceConfigRequest) (*tfprotov6.ValidateResourceConfigResponse, error) {
47-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
49+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
50+
Type: "Resource",
51+
Name: req.TypeName,
52+
Operation: "ValidateResourceConfig",
53+
})
4854
return s.OldServer.ValidateResourceConfig(ctx, req)
4955
}
5056

5157
func (s *WrappedProviderServer) UpgradeResourceState(ctx context.Context, req *tfprotov6.UpgradeResourceStateRequest) (*tfprotov6.UpgradeResourceStateResponse, error) {
52-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, "upgrade."+req.TypeName)
58+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
59+
Type: "Resource",
60+
Name: req.TypeName,
61+
Operation: "UpgradeResourceState",
62+
})
5363
return s.OldServer.UpgradeResourceState(ctx, req)
5464
}
5565

5666
func (s *WrappedProviderServer) ReadResource(ctx context.Context, req *tfprotov6.ReadResourceRequest) (*tfprotov6.ReadResourceResponse, error) {
57-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
67+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
68+
Type: "Resource",
69+
Name: req.TypeName,
70+
Operation: "ReadResource",
71+
})
5872
return s.OldServer.ReadResource(ctx, req)
5973
}
6074

6175
func (s *WrappedProviderServer) PlanResourceChange(ctx context.Context, req *tfprotov6.PlanResourceChangeRequest) (*tfprotov6.PlanResourceChangeResponse, error) {
62-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
76+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
77+
Type: "Resource",
78+
Name: req.TypeName,
79+
Operation: "PlanResourceChange",
80+
})
6381
return s.OldServer.PlanResourceChange(ctx, req)
6482
}
6583

6684
func (s *WrappedProviderServer) ApplyResourceChange(ctx context.Context, req *tfprotov6.ApplyResourceChangeRequest) (*tfprotov6.ApplyResourceChangeResponse, error) {
67-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
85+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
86+
Type: "Resource",
87+
Name: req.TypeName,
88+
Operation: "ApplyResourceChange",
89+
})
6890
return s.OldServer.ApplyResourceChange(ctx, req)
6991
}
7092

7193
func (s *WrappedProviderServer) ImportResourceState(ctx context.Context, req *tfprotov6.ImportResourceStateRequest) (*tfprotov6.ImportResourceStateResponse, error) {
72-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, "import."+req.TypeName)
94+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
95+
Type: "Resource",
96+
Name: req.TypeName,
97+
Operation: "ImportResourceState",
98+
})
7399
return s.OldServer.ImportResourceState(ctx, req)
74100
}
75101

76102
func (s *WrappedProviderServer) MoveResourceState(ctx context.Context, req *tfprotov6.MoveResourceStateRequest) (*tfprotov6.MoveResourceStateResponse, error) {
77-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, "move."+req.TargetTypeName)
103+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
104+
Type: "Resource",
105+
Name: req.TargetTypeName,
106+
Operation: "MoveResourceState",
107+
})
78108
return s.OldServer.MoveResourceState(ctx, req)
79109
}
80110

81111
func (s *WrappedProviderServer) UpgradeResourceIdentity(ctx context.Context, req *tfprotov6.UpgradeResourceIdentityRequest) (*tfprotov6.UpgradeResourceIdentityResponse, error) {
82-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
112+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
113+
Type: "Resource",
114+
Name: req.TypeName,
115+
Operation: "UpgradeResourceIdentity",
116+
})
83117
return s.OldServer.UpgradeResourceIdentity(ctx, req)
84118
}
85119

86120
func (s *WrappedProviderServer) ValidateDataResourceConfig(ctx context.Context, req *tfprotov6.ValidateDataResourceConfigRequest) (*tfprotov6.ValidateDataResourceConfigResponse, error) {
87-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, "data."+req.TypeName)
121+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
122+
Type: "Datasource",
123+
Name: req.TypeName,
124+
Operation: "ValidateDataResourceConfig",
125+
})
88126
return s.OldServer.ValidateDataResourceConfig(ctx, req)
89127
}
90128

91129
func (s *WrappedProviderServer) ReadDataSource(ctx context.Context, req *tfprotov6.ReadDataSourceRequest) (*tfprotov6.ReadDataSourceResponse, error) {
92-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, "data."+req.TypeName)
130+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
131+
Type: "Datasource",
132+
Name: req.TypeName,
133+
Operation: "ReadDataSource",
134+
})
93135
return s.OldServer.ReadDataSource(ctx, req)
94136
}
95137

96138
func (s *WrappedProviderServer) CallFunction(ctx context.Context, req *tfprotov6.CallFunctionRequest) (*tfprotov6.CallFunctionResponse, error) {
97-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, "func."+req.Name)
139+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
140+
Type: "Function",
141+
Name: req.Name,
142+
Operation: "CallFunction",
143+
})
98144
return s.OldServer.CallFunction(ctx, req)
99145
}
100146

@@ -103,21 +149,37 @@ func (s *WrappedProviderServer) GetFunctions(ctx context.Context, req *tfprotov6
103149
}
104150

105151
func (s *WrappedProviderServer) ValidateEphemeralResourceConfig(ctx context.Context, req *tfprotov6.ValidateEphemeralResourceConfigRequest) (*tfprotov6.ValidateEphemeralResourceConfigResponse, error) {
106-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
152+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
153+
Type: "Ephemeral",
154+
Name: req.TypeName,
155+
Operation: "ValidateEphemeralResourceConfig",
156+
})
107157
return s.OldServer.ValidateEphemeralResourceConfig(ctx, req)
108158
}
109159

110160
func (s *WrappedProviderServer) OpenEphemeralResource(ctx context.Context, req *tfprotov6.OpenEphemeralResourceRequest) (*tfprotov6.OpenEphemeralResourceResponse, error) {
111-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
161+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
162+
Type: "Ephemeral",
163+
Name: req.TypeName,
164+
Operation: "OpenEphemeralResource",
165+
})
112166
return s.OldServer.OpenEphemeralResource(ctx, req)
113167
}
114168

115169
func (s *WrappedProviderServer) RenewEphemeralResource(ctx context.Context, req *tfprotov6.RenewEphemeralResourceRequest) (*tfprotov6.RenewEphemeralResourceResponse, error) {
116-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
170+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
171+
Type: "Ephemeral",
172+
Name: req.TypeName,
173+
Operation: "RenewEphemeralResource",
174+
})
117175
return s.OldServer.RenewEphemeralResource(ctx, req)
118176
}
119177

120178
func (s *WrappedProviderServer) CloseEphemeralResource(ctx context.Context, req *tfprotov6.CloseEphemeralResourceRequest) (*tfprotov6.CloseEphemeralResourceResponse, error) {
121-
ctx = context.WithValue(ctx, config.ContextKeyTFSrc, req.TypeName)
179+
ctx = config.AddUserAgentExtra(ctx, config.UserAgentExtra{
180+
Type: "Ephemeral",
181+
Name: req.TypeName,
182+
Operation: "CloseEphemeralResource",
183+
})
122184
return s.OldServer.CloseEphemeralResource(ctx, req)
123185
}

0 commit comments

Comments
 (0)