Skip to content

Commit 4b9df7a

Browse files
stephanosstephanos
committed
Additional env config fields
1 parent 62a0789 commit 4b9df7a

File tree

2 files changed

+220
-28
lines changed

2 files changed

+220
-28
lines changed

‎contrib/envconfig/client_config_toml.go‎

Lines changed: 105 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -8,23 +8,85 @@ import (
88
"github.com/BurntSushi/toml"
99
)
1010

11+
// knownProfileKeys contains the TOML keys recognized for profile configuration.
12+
// This must be kept in sync with tomlClientConfigProfile's TOML tags.
13+
// See TestKnownProfileKeysInSync for validation.
14+
var knownProfileKeys = map[string]bool{
15+
"address": true,
16+
"namespace": true,
17+
"api_key": true,
18+
"tls": true,
19+
"codec": true,
20+
"grpc_meta": true,
21+
}
22+
1123
// ClientConfigToTOMLOptions are options for [ClientConfig.ToTOML].
1224
type ClientConfigToTOMLOptions struct {
1325
// Defaults to two-space indent.
1426
OverrideIndent *string
27+
// If non-nil, these additional fields will be serialized with each profile.
28+
// Key is profile name, value is map of field name to field value.
29+
AdditionalProfileFields map[string]map[string]any
1530
}
1631

1732
// ToTOML converts the client config to TOML. Note, this may not be byte-for-byte exactly what may have been set in
1833
// [ClientConfig.FromTOML].
1934
func (c *ClientConfig) ToTOML(options ClientConfigToTOMLOptions) ([]byte, error) {
2035
var conf tomlClientConfig
2136
conf.fromClientConfig(c)
37+
38+
// If no additional fields, use fast path with direct struct encoding
39+
if len(options.AdditionalProfileFields) == 0 {
40+
var buf bytes.Buffer
41+
enc := toml.NewEncoder(&buf)
42+
if options.OverrideIndent != nil {
43+
enc.Indent = *options.OverrideIndent
44+
}
45+
if err := enc.Encode(&conf); err != nil {
46+
return nil, err
47+
}
48+
return buf.Bytes(), nil
49+
}
50+
51+
// Need to merge additional fields - encode to TOML then decode to map
2252
var buf bytes.Buffer
53+
if err := toml.NewEncoder(&buf).Encode(&conf); err != nil {
54+
return nil, err
55+
}
56+
var rawConf map[string]any
57+
if _, err := toml.Decode(buf.String(), &rawConf); err != nil {
58+
return nil, err
59+
}
60+
61+
// Get or create profile map
62+
profiles, _ := rawConf["profile"].(map[string]any)
63+
if profiles == nil {
64+
profiles = make(map[string]any)
65+
rawConf["profile"] = profiles
66+
}
67+
68+
// Merge additional fields into profiles
69+
for profileName, additional := range options.AdditionalProfileFields {
70+
profile, _ := profiles[profileName].(map[string]any)
71+
if profile == nil {
72+
profile = make(map[string]any)
73+
profiles[profileName] = profile
74+
}
75+
for k, v := range additional {
76+
if knownProfileKeys[k] {
77+
return nil, fmt.Errorf("additional field %q in profile %q conflicts with known profile field", k, profileName)
78+
}
79+
profile[k] = v
80+
}
81+
}
82+
83+
// Re-encode with merged data
84+
buf.Reset()
2385
enc := toml.NewEncoder(&buf)
2486
if options.OverrideIndent != nil {
2587
enc.Indent = *options.OverrideIndent
2688
}
27-
if err := enc.Encode(&conf); err != nil {
89+
if err := enc.Encode(rawConf); err != nil {
2890
return nil, err
2991
}
3092
return buf.Bytes(), nil
@@ -33,24 +95,58 @@ func (c *ClientConfig) ToTOML(options ClientConfigToTOMLOptions) ([]byte, error)
3395
type ClientConfigFromTOMLOptions struct {
3496
// If true, will error if there are unrecognized keys.
3597
Strict bool
98+
// If non-nil, populated with additional (unrecognized) profile fields.
99+
// Key is profile name, value is map of field name to field value.
100+
// This allows callers to preserve custom profile fields without modifying
101+
// this package. Note, if Strict is true the additional fields will cause an
102+
// error before they can be captured here.
103+
AdditionalProfileFields map[string]map[string]any
36104
}
37105

38106
// FromTOML converts from TOML to the client config. This will replace all profiles within, it does not do any form of
39107
// merging.
40108
func (c *ClientConfig) FromTOML(b []byte, options ClientConfigFromTOMLOptions) error {
41109
var conf tomlClientConfig
42-
if md, err := toml.Decode(string(b), &conf); err != nil {
110+
md, err := toml.Decode(string(b), &conf)
111+
if err != nil {
43112
return err
44-
} else if options.Strict {
45-
unknown := md.Undecoded()
46-
if len(unknown) > 0 {
47-
keys := make([]string, len(unknown))
48-
for i, k := range unknown {
49-
keys[i] = k.String()
113+
}
114+
115+
undecoded := md.Undecoded()
116+
if options.Strict && len(undecoded) > 0 {
117+
keys := make([]string, len(undecoded))
118+
for i, k := range undecoded {
119+
keys[i] = k.String()
120+
}
121+
return fmt.Errorf("key(s) unrecognized: %v", strings.Join(keys, ", "))
122+
}
123+
124+
// If AdditionalProfileFields is requested, extract unknown profile fields.
125+
if options.AdditionalProfileFields != nil && len(undecoded) > 0 {
126+
// Decode again into raw map to get additional field values
127+
var rawConf struct {
128+
Profiles map[string]map[string]any `toml:"profile"`
129+
}
130+
if _, err := toml.Decode(string(b), &rawConf); err != nil {
131+
return err
132+
}
133+
134+
for _, key := range undecoded {
135+
// Skip non-profile undecoded keys (e.g., unknown top-level sections)
136+
if len(key) < 3 || key[0] != "profile" {
137+
continue
138+
}
139+
profileName := key[1]
140+
fieldKey := key[2]
141+
if v, ok := rawConf.Profiles[profileName][fieldKey]; ok {
142+
if options.AdditionalProfileFields[profileName] == nil {
143+
options.AdditionalProfileFields[profileName] = make(map[string]any)
144+
}
145+
options.AdditionalProfileFields[profileName][fieldKey] = v
50146
}
51-
return fmt.Errorf("key(s) unrecognized: %v", strings.Join(keys, ", "))
52147
}
53148
}
149+
54150
conf.applyToClientConfig(c)
55151
return nil
56152
}

‎contrib/envconfig/client_config_toml_test.go‎

Lines changed: 115 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1-
package envconfig_test
1+
package envconfig
22

33
import (
4+
"reflect"
5+
"strings"
46
"testing"
57

68
"github.com/stretchr/testify/require"
7-
"go.temporal.io/sdk/contrib/envconfig"
89
)
910

1011
func TestClientConfigTOMLFull(t *testing.T) {
@@ -28,8 +29,8 @@ server_ca_cert_data = "my-server-ca-cert-data"
2829
server_name = "my-server-name"
2930
disable_host_verification = true`
3031

31-
var conf envconfig.ClientConfig
32-
require.NoError(t, conf.FromTOML([]byte(data), envconfig.ClientConfigFromTOMLOptions{}))
32+
var conf ClientConfig
33+
require.NoError(t, conf.FromTOML([]byte(data), ClientConfigFromTOMLOptions{}))
3334
prof := conf.Profiles["foo"]
3435
require.Equal(t, "my-address", prof.Address)
3536
require.Equal(t, "my-namespace", prof.Namespace)
@@ -48,10 +49,10 @@ disable_host_verification = true`
4849
require.Equal(t, map[string]string{"some-header-key": "some-value"}, prof.GRPCMeta)
4950

5051
// Back to toml and back to structure again, then deep equality check
51-
b, err := conf.ToTOML(envconfig.ClientConfigToTOMLOptions{})
52+
b, err := conf.ToTOML(ClientConfigToTOMLOptions{})
5253
require.NoError(t, err)
53-
var newConf envconfig.ClientConfig
54-
require.NoError(t, newConf.FromTOML(b, envconfig.ClientConfigFromTOMLOptions{}))
54+
var newConf ClientConfig
55+
require.NoError(t, newConf.FromTOML(b, ClientConfigFromTOMLOptions{}))
5556
require.Equal(t, conf, newConf)
5657
// Sanity check that require.Equal actually does deep-equality
5758
newConf.Profiles["foo"].Codec.Auth += "-dirty"
@@ -67,8 +68,8 @@ stuff = "does not matter"
6768
address = "my-address"
6869
some_future_key = "some value"`
6970

70-
var conf envconfig.ClientConfig
71-
err := conf.FromTOML([]byte(data), envconfig.ClientConfigFromTOMLOptions{Strict: true})
71+
var conf ClientConfig
72+
err := conf.FromTOML([]byte(data), ClientConfigFromTOMLOptions{Strict: true})
7273
require.ErrorContains(t, err, "unimportant.stuff")
7374
require.ErrorContains(t, err, "profile.foo.some_future_key")
7475
}
@@ -82,8 +83,8 @@ api_key = "my-api-key"
8283
[profile.foo.tls]
8384
`
8485

85-
var conf envconfig.ClientConfig
86-
require.NoError(t, conf.FromTOML([]byte(data), envconfig.ClientConfigFromTOMLOptions{}))
86+
var conf ClientConfig
87+
require.NoError(t, conf.FromTOML([]byte(data), ClientConfigFromTOMLOptions{}))
8788
prof := conf.Profiles["foo"]
8889
require.Empty(t, prof.Address)
8990
require.Empty(t, prof.Namespace)
@@ -93,22 +94,117 @@ api_key = "my-api-key"
9394
require.Zero(t, *prof.TLS)
9495

9596
// Back to toml and back to structure again, then deep equality check
96-
b, err := conf.ToTOML(envconfig.ClientConfigToTOMLOptions{})
97+
b, err := conf.ToTOML(ClientConfigToTOMLOptions{})
9798
require.NoError(t, err)
98-
var newConf envconfig.ClientConfig
99-
require.NoError(t, newConf.FromTOML(b, envconfig.ClientConfigFromTOMLOptions{}))
99+
var newConf ClientConfig
100+
require.NoError(t, newConf.FromTOML(b, ClientConfigFromTOMLOptions{}))
100101
require.Equal(t, conf, newConf)
101102
}
102103

103104
func TestClientConfigTOMLEmpty(t *testing.T) {
104-
var conf envconfig.ClientConfig
105-
require.NoError(t, conf.FromTOML(nil, envconfig.ClientConfigFromTOMLOptions{}))
105+
var conf ClientConfig
106+
require.NoError(t, conf.FromTOML(nil, ClientConfigFromTOMLOptions{}))
106107
require.Empty(t, conf.Profiles)
107108

108109
// Back to toml and back to structure again, then deep equality check
109-
b, err := conf.ToTOML(envconfig.ClientConfigToTOMLOptions{})
110+
b, err := conf.ToTOML(ClientConfigToTOMLOptions{})
110111
require.NoError(t, err)
111-
var newConf envconfig.ClientConfig
112-
require.NoError(t, newConf.FromTOML(b, envconfig.ClientConfigFromTOMLOptions{}))
112+
var newConf ClientConfig
113+
require.NoError(t, newConf.FromTOML(b, ClientConfigFromTOMLOptions{}))
113114
require.Equal(t, conf, newConf)
114115
}
116+
117+
func TestClientConfigTOMLAdditionalProfileFields(t *testing.T) {
118+
data := `
119+
[profile.foo]
120+
address = "my-address"
121+
namespace = "my-namespace"
122+
custom_field = "custom-value"
123+
custom_field2 = 42
124+
125+
[profile.foo.custom_nested]
126+
key1 = "value1"
127+
128+
[profile.foo.custom_nested.deep]
129+
key2 = "value2"
130+
131+
[profile.foo.custom_nested.deep.deeper]
132+
key3 = "value3"
133+
134+
[profile.bar]
135+
address = "bar-address"
136+
custom_field = true`
137+
138+
var conf ClientConfig
139+
additional := make(map[string]map[string]any)
140+
require.NoError(t, conf.FromTOML([]byte(data), ClientConfigFromTOMLOptions{
141+
AdditionalProfileFields: additional,
142+
}))
143+
144+
// Verify known fields were parsed
145+
require.Equal(t, "my-address", conf.Profiles["foo"].Address)
146+
require.Equal(t, "my-namespace", conf.Profiles["foo"].Namespace)
147+
require.Equal(t, "bar-address", conf.Profiles["bar"].Address)
148+
149+
// Verify additional fields were captured
150+
require.Equal(t, "custom-value", additional["foo"]["custom_field"])
151+
require.Equal(t, int64(42), additional["foo"]["custom_field2"])
152+
require.Equal(t, true, additional["bar"]["custom_field"])
153+
154+
// Verify deeply nested additional fields are preserved
155+
customNested, ok := additional["foo"]["custom_nested"].(map[string]any)
156+
require.True(t, ok, "custom_nested should be a map")
157+
require.Equal(t, "value1", customNested["key1"])
158+
159+
deep, ok := customNested["deep"].(map[string]any)
160+
require.True(t, ok, "custom_nested.deep should be a map")
161+
require.Equal(t, "value2", deep["key2"])
162+
163+
deeper, ok := deep["deeper"].(map[string]any)
164+
require.True(t, ok, "custom_nested.deep.deeper should be a map")
165+
require.Equal(t, "value3", deeper["key3"])
166+
167+
// Back to TOML and back to structure again, then deep equality check
168+
b, err := conf.ToTOML(ClientConfigToTOMLOptions{
169+
AdditionalProfileFields: additional,
170+
})
171+
require.NoError(t, err)
172+
var newConf ClientConfig
173+
newAdditional := make(map[string]map[string]any)
174+
require.NoError(t, newConf.FromTOML(b, ClientConfigFromTOMLOptions{
175+
AdditionalProfileFields: newAdditional,
176+
}))
177+
require.Equal(t, conf, newConf)
178+
require.Equal(t, additional, newAdditional)
179+
}
180+
181+
func TestClientConfigTOMLAdditionalProfileFieldsConflict(t *testing.T) {
182+
conf := ClientConfig{
183+
Profiles: map[string]*ClientConfigProfile{
184+
"foo": {Address: "my-address"},
185+
},
186+
}
187+
188+
// Attempt to write with an additional field that conflicts with a known field
189+
_, err := conf.ToTOML(ClientConfigToTOMLOptions{
190+
AdditionalProfileFields: map[string]map[string]any{
191+
"foo": {"address": "conflict"},
192+
},
193+
})
194+
require.ErrorContains(t, err, "additional field \"address\" in profile \"foo\" conflicts with known profile field")
195+
}
196+
197+
func TestKnownProfileKeysInSync(t *testing.T) {
198+
// Extract keys from tomlClientConfigProfile's TOML tags using reflection
199+
expected := make(map[string]bool)
200+
typ := reflect.TypeFor[tomlClientConfigProfile]()
201+
for i := range typ.NumField() {
202+
tag := typ.Field(i).Tag.Get("toml")
203+
if key, _, _ := strings.Cut(tag, ","); key != "" {
204+
expected[key] = true
205+
}
206+
}
207+
208+
require.Equal(t, expected, knownProfileKeys,
209+
"knownProfileKeys must match tomlClientConfigProfile's TOML tags")
210+
}

0 commit comments

Comments
 (0)