Skip to content

Commit cea1325

Browse files
yanand0909yanand0909
committed
[MATRIX-881] ADDED Dynamic required field validation
1 parent fe0636b commit cea1325

File tree

2 files changed

+173
-46
lines changed

2 files changed

+173
-46
lines changed

‎internal/flink/command_connection.go‎

Lines changed: 122 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,12 @@ func AddConnectionSecretFlags(cmd *cobra.Command) {
8787
cmd.Flags().String("service-key", "", fmt.Sprintf("Specify service key for the type: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["service-key"], "or")))
8888
cmd.Flags().String("username", "", fmt.Sprintf("Specify username for the type: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["username"], "or")))
8989
cmd.Flags().String("password", "", fmt.Sprintf("Specify password for the type: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["password"], "or")))
90+
cmd.Flags().String("auth-type", "", fmt.Sprintf("Specify authentication type for the type : %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["auth-type"], "or")))
91+
cmd.Flags().String("bearer-token", "", fmt.Sprintf("Specify bearer token for BEARER authentication type: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["bearer-token"], "or")))
92+
cmd.Flags().String("oauth2-token-endpoint", "", fmt.Sprintf("Specify oauth2 token endpoint: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["oauth2-token-endpoint"], "or")))
93+
cmd.Flags().String("oauth2-client-id", "", fmt.Sprintf("Specify oauth2 client id: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["oauth2-client-id"], "or")))
94+
cmd.Flags().String("oauth2-client-secret", "", fmt.Sprintf("Specify oauth2 client secret: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["oauth2-client-secret"], "or")))
95+
cmd.Flags().String("oauth2-scope", "", fmt.Sprintf("Specify oauth2 scope: %s.", utils.ArrayToCommaDelimitedString(flink.ConnectionSecretTypeMapping["oauth2-scope"], "or")))
9096
}
9197

9298
func validateConnectionType(connectionType string) error {
@@ -100,54 +106,139 @@ func validateConnectionSecrets(cmd *cobra.Command, connectionType string) (map[s
100106
var connectionSecrets []string
101107
connectionSecrets = append(connectionSecrets, flink.ConnectionTypeSecretMapping[connectionType]...)
102108

103-
for key := range flink.ConnectionSecretTypeMapping {
104-
secret, err := cmd.Flags().GetString(key)
109+
secretMap := map[string]string{}
110+
var requiredSecretKeys []string
111+
var optionalSecretKeys []string
112+
113+
dynamicKey, hasDynamicKey := flink.ConnectionTypeDynamicKeyMapping[connectionType]
114+
if hasDynamicKey {
115+
dynamicKeyValue, err := cmd.Flags().GetString(dynamicKey)
105116
if err != nil {
106117
return nil, err
107118
}
108-
if secret != "" && !slices.Contains(connectionSecrets, key) {
109-
return nil, errors.NewErrorWithSuggestions(fmt.Sprintf("%s is invalid for connection %s.", key, connectionType), fmt.Sprintf("Valid secret types are %s.", utils.ArrayToCommaDelimitedString(connectionSecrets, "or")))
119+
if dynamicKeyValue == "" {
120+
return nil, fmt.Errorf("must provide %s for connection %s", dynamicKey, connectionType)
110121
}
111-
}
112122

113-
requiredSecretKeys := flink.ConnectionRequiredSecretMapping[connectionType]
114-
var optionalSecretKeys []string
115-
for _, secretKey := range flink.ConnectionTypeSecretMapping[connectionType] {
116-
if !slices.Contains(requiredSecretKeys, secretKey) {
117-
optionalSecretKeys = append(optionalSecretKeys, secretKey)
123+
backendKey, ok := flink.ConnectionSecretBackendKeyMapping[dynamicKey]
124+
if !ok {
125+
return nil, fmt.Errorf(`backend key not found for "%s"`, dynamicKey)
126+
}
127+
secretMap[backendKey] = dynamicKeyValue
128+
129+
requiredSecretKeys, exists := flink.ConnectionDynamicRequiredSecretMapping[connectionType][dynamicKeyValue]
130+
if !exists {
131+
validTypes := make([]string, 0, len(flink.ConnectionDynamicRequiredSecretMapping[connectionType]))
132+
for k := range flink.ConnectionDynamicRequiredSecretMapping[connectionType] {
133+
validTypes = append(validTypes, k)
134+
}
135+
return nil, errors.NewErrorWithSuggestions(
136+
fmt.Sprintf("invalid %s %s for connection %s", dynamicKey, dynamicKeyValue, connectionType),
137+
fmt.Sprintf("Valid types are %s.", utils.ArrayToCommaDelimitedString(validTypes, "or")),
138+
)
118139
}
119-
}
120140

121-
secretMap := map[string]string{}
122-
for _, requiredKey := range requiredSecretKeys {
123-
secret, err := cmd.Flags().GetString(requiredKey)
124-
if err != nil {
125-
return nil, err
141+
allPossibleKeys, exists := flink.ConnectionDynamicSecretMapping[connectionType][dynamicKeyValue]
142+
if exists {
143+
connectionSecrets = append(connectionSecrets, allPossibleKeys...)
144+
}
145+
146+
for key := range flink.ConnectionSecretTypeMapping {
147+
secret, err := cmd.Flags().GetString(key)
148+
if err != nil {
149+
return nil, err
150+
}
151+
if secret != "" && !slices.Contains(connectionSecrets, key) {
152+
return nil, errors.NewErrorWithSuggestions(
153+
fmt.Sprintf("%s is invalid for connection %s with %s %s.", key, connectionType, dynamicKey, dynamicKeyValue),
154+
fmt.Sprintf("Valid secret types are %s.", utils.ArrayToCommaDelimitedString(connectionSecrets, "or")),
155+
)
156+
}
126157
}
127-
if secret == "" {
128-
return nil, fmt.Errorf("must provide %s for type %s", requiredKey, connectionType)
158+
159+
for _, secretKey := range allPossibleKeys {
160+
if !slices.Contains(requiredSecretKeys, secretKey) {
161+
optionalSecretKeys = append(optionalSecretKeys, secretKey)
162+
}
129163
}
130-
backendKey, ok := flink.ConnectionSecretBackendKeyMapping[requiredKey]
131-
if !ok {
132-
return nil, fmt.Errorf(`backend key not found for "%s"`, requiredKey)
164+
165+
for _, requiredKey := range requiredSecretKeys {
166+
secret, err := cmd.Flags().GetString(requiredKey)
167+
if err != nil {
168+
return nil, err
169+
}
170+
if secret == "" {
171+
return nil, fmt.Errorf("must provide %s for %s %s on connection %s", requiredKey, dynamicKey, dynamicKeyValue, connectionType)
172+
}
173+
backendKey, ok := flink.ConnectionSecretBackendKeyMapping[requiredKey]
174+
if !ok {
175+
return nil, fmt.Errorf(`backend key not found for "%s"`, requiredKey)
176+
}
177+
secretMap[backendKey] = secret
133178
}
134-
secretMap[backendKey] = secret
135-
}
136179

137-
for _, optionalSecretKey := range optionalSecretKeys {
138-
secret, err := cmd.Flags().GetString(optionalSecretKey)
139-
if err != nil {
140-
return nil, err
180+
for _, optionalKey := range optionalSecretKeys {
181+
secret, err := cmd.Flags().GetString(optionalKey)
182+
if err != nil {
183+
return nil, err
184+
}
185+
if secret != "" {
186+
backendKey, ok := flink.ConnectionSecretBackendKeyMapping[optionalKey]
187+
if !ok {
188+
return nil, fmt.Errorf(`backend key not found for "%s"`, optionalKey)
189+
}
190+
secretMap[backendKey] = secret
191+
}
192+
}
193+
} else {
194+
for key := range flink.ConnectionSecretTypeMapping {
195+
secret, err := cmd.Flags().GetString(key)
196+
if err != nil {
197+
return nil, err
198+
}
199+
if secret != "" && !slices.Contains(connectionSecrets, key) {
200+
return nil, errors.NewErrorWithSuggestions(
201+
fmt.Sprintf("%s is invalid for connection %s.", key, connectionType),
202+
fmt.Sprintf("Valid secret types are %s.", utils.ArrayToCommaDelimitedString(connectionSecrets, "or")),
203+
)
204+
}
141205
}
142206

143-
backendKey, ok := flink.ConnectionSecretBackendKeyMapping[optionalSecretKey]
144-
if !ok {
145-
return nil, fmt.Errorf("backend key not found for %s", optionalSecretKey)
207+
requiredSecretKeys = flink.ConnectionRequiredSecretMapping[connectionType]
208+
for _, secretKey := range flink.ConnectionTypeSecretMapping[connectionType] {
209+
if !slices.Contains(requiredSecretKeys, secretKey) {
210+
optionalSecretKeys = append(optionalSecretKeys, secretKey)
211+
}
146212
}
147213

148-
if secret != "" {
214+
for _, requiredKey := range requiredSecretKeys {
215+
secret, err := cmd.Flags().GetString(requiredKey)
216+
if err != nil {
217+
return nil, err
218+
}
219+
if secret == "" {
220+
return nil, fmt.Errorf("must provide %s for type %s", requiredKey, connectionType)
221+
}
222+
backendKey, ok := flink.ConnectionSecretBackendKeyMapping[requiredKey]
223+
if !ok {
224+
return nil, fmt.Errorf(`backend key not found for "%s"`, requiredKey)
225+
}
149226
secretMap[backendKey] = secret
150227
}
228+
229+
for _, optionalKey := range optionalSecretKeys {
230+
secret, err := cmd.Flags().GetString(optionalKey)
231+
if err != nil {
232+
return nil, err
233+
}
234+
if secret != "" {
235+
backendKey, ok := flink.ConnectionSecretBackendKeyMapping[optionalKey]
236+
if !ok {
237+
return nil, fmt.Errorf("backend key not found for %s", optionalKey)
238+
}
239+
secretMap[backendKey] = secret
240+
}
241+
}
151242
}
152243

153244
return secretMap, nil

‎pkg/flink/utils.go‎

Lines changed: 51 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package flink
22

33
var (
4-
ConnectionTypes = []string{"openai", "azureml", "azureopenai", "bedrock", "sagemaker", "googleai", "vertexai", "mongodb", "elastic", "pinecone", "couchbase", "confluent_jdbc"}
4+
ConnectionTypes = []string{"openai", "azureml", "azureopenai", "bedrock", "sagemaker", "googleai", "vertexai", "mongodb", "elastic", "pinecone", "couchbase", "confluent_jdbc", "mcp_server"}
55
ConnectionTypeSecretMapping = map[string][]string{
66
"openai": {"api-key"},
77
"azureml": {"api-key"},
@@ -15,16 +15,23 @@ var (
1515
"pinecone": {"api-key"},
1616
"couchbase": {"username", "password"},
1717
"confluent_jdbc": {"username", "password"},
18+
"mcp_server": {"auth-type"},
1819
}
1920

2021
ConnectionSecretTypeMapping = map[string][]string{
21-
"api-key": {"openai", "azureml", "azureopenai", "googleai", "elastic", "pinecone"},
22-
"aws-access-key": {"bedrock", "sagemaker"},
23-
"aws-secret-key": {"bedrock", "sagemaker"},
24-
"aws-session-token": {"bedrock", "sagemaker"},
25-
"service-key": {"vertexai"},
26-
"username": {"mongodb", "couchbase", "confluent_jdbc"},
27-
"password": {"mongodb", "couchbase", "confluent_jdbc"},
22+
"api-key": {"openai", "azureml", "azureopenai", "googleai", "elastic", "pinecone", "mcp_server"},
23+
"aws-access-key": {"bedrock", "sagemaker"},
24+
"aws-secret-key": {"bedrock", "sagemaker"},
25+
"aws-session-token": {"bedrock", "sagemaker"},
26+
"service-key": {"vertexai"},
27+
"username": {"mongodb", "couchbase", "confluent_jdbc"},
28+
"password": {"mongodb", "couchbase", "confluent_jdbc"},
29+
"auth-type": {"mcp_server"},
30+
"bearer-token": {"mcp_server"},
31+
"oauth2-token-endpoint": {"mcp_server"},
32+
"oauth2-client-secret": {"mcp_server"},
33+
"oauth2-client-id": {"mcp_server"},
34+
"oauth2-scope": {"mcp_server"},
2835
}
2936

3037
ConnectionRequiredSecretMapping = map[string][]string{
@@ -40,14 +47,43 @@ var (
4047
"pinecone": {"api-key"},
4148
"couchbase": {"username", "password"},
4249
"confluent_jdbc": {"username", "password"},
50+
"mcp_server": {"auth-type"},
4351
}
4452
ConnectionSecretBackendKeyMapping = map[string]string{
45-
"api-key": "API_KEY",
46-
"aws-access-key": "AWS_ACCESS_KEY_ID",
47-
"aws-secret-key": "AWS_SECRET_ACCESS_KEY",
48-
"aws-session-token": "AWS_SESSION_TOKEN",
49-
"service-key": "SERVICE_KEY",
50-
"username": "USERNAME",
51-
"password": "PASSWORD",
53+
"api-key": "API_KEY",
54+
"aws-access-key": "AWS_ACCESS_KEY_ID",
55+
"aws-secret-key": "AWS_SECRET_ACCESS_KEY",
56+
"aws-session-token": "AWS_SESSION_TOKEN",
57+
"service-key": "SERVICE_KEY",
58+
"username": "USERNAME",
59+
"password": "PASSWORD",
60+
"auth-type": "AUTH_TYPE",
61+
"bearer-token": "BEARER_TOKEN",
62+
"oauth2-token-endpoint": "OAUTH2_TOKEN_ENDPOINT",
63+
"oauth2-client-secret": "OAUTH2_CLIENT_SECRET",
64+
"oauth2-client-id": "OAUTH2_CLIENT_ID",
65+
"oauth2-scope": "OAUTH2_SCOPE",
66+
}
67+
68+
ConnectionTypeDynamicKeyMapping = map[string]string{
69+
"mcp_server": "auth-type",
70+
}
71+
72+
ConnectionDynamicRequiredSecretMapping = map[string]map[string][]string{
73+
"mcp_server": {
74+
"NO_AUTH": {},
75+
"API_KEY": {"api-key"},
76+
"BEARER": {"bearer-token"},
77+
"OAUTH2": {"oauth2-token-endpoint", "oauth2-client-id", "oauth2-client-secret", "oauth2-scope"},
78+
},
79+
}
80+
81+
ConnectionDynamicSecretMapping = map[string]map[string][]string{
82+
"mcp_server": {
83+
"NO_AUTH": {},
84+
"API_KEY": {"api-key"},
85+
"BEARER": {"bearer-token"},
86+
"OAUTH2": {"oauth2-token-endpoint", "oauth2-client-id", "oauth2-client-secret", "oauth2-scope"},
87+
},
5288
}
5389
)

0 commit comments

Comments
 (0)