Repository: ariga/atlas
Pull request #1217 (2022-10-20 16:23:17)

schemahcl: added `urlqueryset` function
Signed-off-by: Giau. Tran Minh <hello@giautm.dev>


Patch:
diff --git a/schemahcl/stdlib.go b/schemahcl/stdlib.go
index 4614e75e..0064f4ae 100644
--- a/schemahcl/stdlib.go
+++ b/schemahcl/stdlib.go
@@ -156,6 +156,7 @@ func stdFuncs() map[string]function.Function {
 		"trimsuffix":      stdlib.TrimSuffixFunc,
 		"try":             tryfunc.TryFunc,
 		"upper":           stdlib.UpperFunc,
+		"urlqueryset":     urlQuerySetFunc,
 		"urlsetpath":      urlSetPathFunc,
 		"values":          stdlib.ValuesFunc,
 		"zipmap":          stdlib.ZipmapFunc,
@@ -238,6 +239,34 @@ func makeToFunc(wantTy cty.Type) function.Function {
 	})
 }
 
+var urlQuerySetFunc = function.New(&function.Spec{
+	Params: []function.Parameter{
+		{
+			Name: "url",
+			Type: cty.String,
+		},
+		{
+			Name: "key",
+			Type: cty.String,
+		},
+		{
+			Name: "value",
+			Type: cty.String,
+		},
+	},
+	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
+		}
+		q := u.Query()
+		q.Set(args[1].AsString(), args[2].AsString())
+		u.RawQuery = q.Encode()
+		return cty.StringVal(u.String()), nil
+	},
+})
+
 var urlSetPathFunc = function.New(&function.Spec{
 	Params: []function.Parameter{
 		{


Test patch:
diff --git a/internal/integration/testdata/postgres/cli-migrate-apply-datasrc.txt b/internal/integration/testdata/postgres/cli-migrate-apply-datasrc.txt
index 859c0480..887290c9 100644
--- a/internal/integration/testdata/postgres/cli-migrate-apply-datasrc.txt
+++ b/internal/integration/testdata/postgres/cli-migrate-apply-datasrc.txt
@@ -13,11 +13,6 @@ variable "pattern" {
   type = string
 }
 
-locals {
-  # The database URL without its schema name.
-  base_url = regexreplace(var.url, "&search_path=\\w+", "")
-}
-
 data "sql" "tenants" {
   url = var.url
   query = <<EOS
@@ -30,7 +25,7 @@ EOS
 
 env "dev" {
   for_each = toset(data.sql.tenants.values)
-  url = "${local.base_url}&search_path=${each.value}"
+  url      = urlqueryset(var.url, "search_path", each.value)
 }
 
 -- migrations/1_create_users.sql --
diff --git a/schemahcl/stdlib_test.go b/schemahcl/stdlib_test.go
index 361ed4fb..0baaaa02 100644
--- a/schemahcl/stdlib_test.go
+++ b/schemahcl/stdlib_test.go
@@ -47,3 +47,49 @@ func TestURLSetPathFunc(t *testing.T) {
 		})
 	}
 }
+
+func TestURLQuerySetFunc(t *testing.T) {
+	tests := []struct {
+		URL   cty.Value
+		Key   cty.Value
+		Value cty.Value
+		Want  cty.Value
+	}{
+		{
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database?sslmode=disable&sslmode=disable"),
+			cty.StringVal("search_path"),
+			cty.StringVal("schema"),
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database?search_path=schema&sslmode=disable&sslmode=disable"),
+		},
+		{
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database?sslmode=disable&search_path=admin&sslmode=disable"),
+			cty.StringVal("search_path"),
+			cty.StringVal("schema"),
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database?search_path=schema&sslmode=disable&sslmode=disable"),
+		},
+		{
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database"),
+			cty.StringVal("search_path"),
+			cty.StringVal("schema"),
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database?search_path=schema"),
+		},
+		{
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database?sslmode=disable&search_path=admin&sslmode=disable"),
+			cty.StringVal("search_path"),
+			cty.StringVal(""),
+			cty.StringVal("postgres://postgres:pass@0.0.0.0:5432/database?search_path=&sslmode=disable&sslmode=disable"),
+		},
+	}
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("(%#v,%#v,%#v)", test.URL, test.Key, test.Value), func(t *testing.T) {
+			got, err := urlQuerySetFunc.Call([]cty.Value{test.URL, test.Key, test.Value})
+
+			if err != nil {
+				t.Fatalf("unexpected error: %s", err)
+			}
+			if !got.RawEquals(test.Want) {
+				t.Errorf("wrong result\ngot:  %#v\nwant: %#v", got, test.Want)
+			}
+		})
+	}
+}