Repository: xataio/pgstream
Pull request #340 (2025-05-14 18:18:55)

Template transformer: Add greenmask toolkit functions
This PR adds support for greenmask's functions for templates.
https://github.com/eminano/greenmask/blob/main/pkg/toolkit/template_functions.go#L71-L199



github-actions[bot]: ### Merging this branch will **increase** overall coverage

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/pkg/transformers | 85.83% (**+0.63%**) | :thumbsup: |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/pkg/transformers/template_transformer.go | 82.35% (**+82.35%**) | 17 (+17) | 14 (+14) | 3 (+3) | :star2: |

_Please note that the "Total", "Covered", and "Missed" counts above refer to ***code statements*** instead of lines of code. The value in brackets refers to the test coverage of that file in the old version of the code._

### Changed unit test files

- github.com/xataio/pgstream/pkg/transformers/template_transformer_test.go

</details>


Patch:
diff --git a/pkg/transformers/template_transformer.go b/pkg/transformers/template_transformer.go
index 84d2bf4..fc4bb28 100644
--- a/pkg/transformers/template_transformer.go
+++ b/pkg/transformers/template_transformer.go
@@ -8,6 +8,8 @@ import (
 	"fmt"
 	"strings"
 	"text/template"
+
+	greenmasktoolkit "github.com/eminano/greenmask/pkg/toolkit"
 )
 
 type TemplateTransformer struct {
@@ -31,7 +33,7 @@ func NewTemplateTransformer(params Parameters) (*TemplateTransformer, error) {
 		return nil, errTemplateMustBeProvided
 	}
 
-	tmpl, err := template.New("").Parse(templateStr)
+	tmpl, err := template.New("").Funcs(greenmasktoolkit.FuncMap()).Parse(templateStr)
 	if err != nil {
 		return nil, fmt.Errorf("template_transformer: error parsing template: %w", err)
 	}


Test patch:
diff --git a/pkg/transformers/template_transformer_test.go b/pkg/transformers/template_transformer_test.go
index 3e9b96c..432190d 100644
--- a/pkg/transformers/template_transformer_test.go
+++ b/pkg/transformers/template_transformer_test.go
@@ -221,3 +221,61 @@ func TestTemplateTransformer_Transform_WithDynamicValues(t *testing.T) {
 		})
 	}
 }
+
+func TestTemplateTransformer_Transform_WithGreenmaskToolkitFuncs(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name          string
+		value         any
+		dynamicValues map[string]any
+		params        Parameters
+		wantOutput    string
+		wantErr       error
+	}{
+		{
+			name:  "ok - dynamic values with masking func",
+			value: "email",
+			dynamicValues: map[string]any{
+				"value1": nil,
+				"value2": "john.doe@xata.io",
+			},
+			params: Parameters{
+				"template": "{{ $first := .GetDynamicValue \"value1\" }}{{ $second :=.GetDynamicValue \"value2\" }} {{- if eq $first nil -}} {{ masking .GetValue $second }} {{- else -}} {{ masking .GetValue $first }} {{- end -}}",
+			},
+			wantOutput: "joh****e@xata.io",
+			wantErr:    nil,
+		},
+		{
+			name:  "ok - random integer",
+			value: 3,
+			params: Parameters{
+				"template": "{{ $randval := randomInt 0 .GetValue}} {{- if and (isInt $randval) (gt $randval 0) (lt $randval .GetValue) -}} {{\"yes\"}} {{- else -}} {{\"no\"}} {{- end -}}",
+			},
+			wantOutput: "yes",
+			wantErr:    nil,
+		},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			tt, err := NewTemplateTransformer(tc.params)
+			require.NoError(t, err)
+			got, err := tt.Transform(context.Background(), Value{
+				TransformValue: tc.value,
+				DynamicValues:  tc.dynamicValues,
+			})
+
+			if tc.wantErr != nil {
+				require.Error(t, err)
+				if !errors.Is(err, tc.wantErr) {
+					require.Contains(t, err.Error(), tc.wantErr.Error())
+				}
+				return
+			}
+			require.NoError(t, err)
+
+			require.Equal(t, tc.wantOutput, got)
+		})
+	}
+}