Repository: ariga/atlas
Pull request #2073 (2023-09-12 15:17:37)

schemahcl: add `urluserinfo` function


Patch:
diff --git a/schemahcl/stdlib.go b/schemahcl/stdlib.go
index 8c60a130..971205be 100644
--- a/schemahcl/stdlib.go
+++ b/schemahcl/stdlib.go
@@ -161,7 +161,8 @@ func stdFuncs() map[string]function.Function {
 		"trimsuffix":      stdlib.TrimSuffixFunc,
 		"try":             tryfunc.TryFunc,
 		"upper":           stdlib.UpperFunc,
-		"urlescape":       urlEscape,
+		"urlescape":       urlEscapeFunc,
+		"urluserinfo":     urlUserinfoFunc,
 		"urlqueryset":     urlQuerySetFunc,
 		"urlsetpath":      urlSetPathFunc,
 		"values":          stdlib.ValuesFunc,
@@ -276,6 +277,37 @@ var (
 			return cty.StringVal(u.String()), nil
 		},
 	})
+	urlUserinfoFunc = function.New(&function.Spec{
+		Params: []function.Parameter{
+			{
+				Name: "url",
+				Type: cty.String,
+			},
+			{
+				Name: "user",
+				Type: cty.String,
+			},
+		},
+		VarParam: &function.Parameter{
+			Name:      "pass",
+			Type:      cty.String,
+			AllowNull: true,
+		},
+		Type: function.StaticReturnType(cty.String),
+		Impl: func(args []cty.Value, retType cty.Type) (cty.Value, error) {
+			u, err := url.Parse(args[0].AsString())
+			if err != nil {
+				return cty.NilVal, err
+			}
+			user := args[1].AsString()
+			if len(args) > 2 && !args[2].IsNull() {
+				u.User = url.UserPassword(user, args[2].AsString())
+			} else {
+				u.User = url.User(user)
+			}
+			return cty.StringVal(u.String()), nil
+		},
+	})
 	urlSetPathFunc = function.New(&function.Spec{
 		Params: []function.Parameter{
 			{
@@ -298,7 +330,7 @@ var (
 		},
 	})
 
-	urlEscape = function.New(&function.Spec{
+	urlEscapeFunc = function.New(&function.Spec{
 		Params: []function.Parameter{
 			{
 				Name: "string",


Test patch:
diff --git a/schemahcl/stdlib_test.go b/schemahcl/stdlib_test.go
index c7933d4a..cdb9fc58 100644
--- a/schemahcl/stdlib_test.go
+++ b/schemahcl/stdlib_test.go
@@ -111,7 +111,7 @@ func TestURLQuerySetFunc(t *testing.T) {
 func TestURLEscapeFunc(t *testing.T) {
 	for _, tt := range []string{"foo", "foo?", "foo&"} {
 		t.Run(tt, func(t *testing.T) {
-			got, err := urlEscape.Call([]cty.Value{cty.StringVal(tt)})
+			got, err := urlEscapeFunc.Call([]cty.Value{cty.StringVal(tt)})
 			if err != nil {
 				t.Fatalf("unexpected error: %s", err)
 			}
@@ -122,6 +122,22 @@ func TestURLEscapeFunc(t *testing.T) {
 	}
 }
 
+func TestURLUserinfoFunc(t *testing.T) {
+	u := cty.StringVal("mysql://localhost:3306")
+	// Only user is provided
+	got, err := urlUserinfoFunc.Call([]cty.Value{u, cty.StringVal("user")})
+	require.NoError(t, err)
+	require.Equal(t, cty.StringVal("mysql://user@localhost:3306"), got)
+	// The password is null
+	got, err = urlUserinfoFunc.Call([]cty.Value{u, cty.StringVal("user"), cty.NullVal(cty.String)})
+	require.NoError(t, err)
+	require.Equal(t, cty.StringVal("mysql://user@localhost:3306"), got)
+	// Both user and password are provided
+	got, err = urlUserinfoFunc.Call([]cty.Value{u, cty.StringVal("user"), cty.StringVal("pass")})
+	require.NoError(t, err)
+	require.Equal(t, cty.StringVal("mysql://user:pass@localhost:3306"), got)
+}
+
 func TestMakeFileFunc(t *testing.T) {
 	fn := MakeFileFunc("testdata")
 	_, err := fn.Call([]cty.Value{cty.StringVal("foo")})