Repository: xataio/pgstream
Pull request #160 (2025-03-03 14:35:00)

Add choice gm transformer
Introduce `ChoiceTransformer` which uses greenmask's `RandomChoiceTransformer` under the hood.

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

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/pkg/transformers | 100.00% (ø) |  |
| github.com/xataio/pgstream/pkg/transformers/builder | 11.11% (**-1.01%**) | :thumbsdown: |
| github.com/xataio/pgstream/pkg/transformers/greenmask | 94.65% (**-0.56%**) | :thumbsdown: |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/pkg/transformers/builder/transformer_builder.go | 11.11% (**-1.01%**) | 36 (+3) | 4 | 32 (+3) | :thumbsdown: |
| github.com/xataio/pgstream/pkg/transformers/greenmask/greenmask_choice_transformer.go | 90.00% (**+90.00%**) | 20 (+20) | 18 (+18) | 2 (+2) | :star2: |
| github.com/xataio/pgstream/pkg/transformers/transformer.go | 100.00% (ø) | 7 | 7 | 0 |  |

_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/greenmask/greenmask_choice_transformer_test.go

</details>

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

| Impacted Packages | Coverage Δ | :robot: |
|-------------------|------------|---------|
| github.com/xataio/pgstream/pkg/transformers | 100.00% (ø) |  |
| github.com/xataio/pgstream/pkg/transformers/builder | 11.11% (**-1.01%**) | :thumbsdown: |
| github.com/xataio/pgstream/pkg/transformers/greenmask | 94.74% (**-0.47%**) | :thumbsdown: |

---

<details>

<summary>Coverage by file</summary>

### Changed files (no unit tests)

| Changed File | Coverage Δ | Total | Covered | Missed | :robot: |
|--------------|------------|-------|---------|--------|---------|
| github.com/xataio/pgstream/pkg/transformers/builder/transformer_builder.go | 11.11% (**-1.01%**) | 36 (+3) | 4 | 32 (+3) | :thumbsdown: |
| github.com/xataio/pgstream/pkg/transformers/greenmask/greenmask_choice_transformer.go | 91.30% (**+91.30%**) | 23 (+23) | 21 (+21) | 2 (+2) | :star2: |
| github.com/xataio/pgstream/pkg/transformers/transformer.go | 100.00% (ø) | 7 | 7 | 0 |  |

_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/greenmask/greenmask_choice_transformer_test.go

</details>


Patch:
diff --git a/pkg/transformers/builder/transformer_builder.go b/pkg/transformers/builder/transformer_builder.go
index 45ddda1..6265100 100644
--- a/pkg/transformers/builder/transformer_builder.go
+++ b/pkg/transformers/builder/transformer_builder.go
@@ -28,6 +28,8 @@ func New(cfg *transformers.Config) (transformers.Transformer, error) {
 		return greenmask.NewUUIDTransformer(cfg.Generator)
 	case transformers.GreenmaskBoolean:
 		return greenmask.NewBooleanTransformer(cfg.Generator)
+	case transformers.GreenmaskChoice:
+		return greenmask.NewChoiceTransformer(cfg.Generator, cfg.Parameters)
 	default:
 		return nil, transformers.ErrUnsupportedTransformer
 	}
diff --git a/pkg/transformers/greenmask/greenmask_choice_transformer.go b/pkg/transformers/greenmask/greenmask_choice_transformer.go
new file mode 100644
index 0000000..af2c908
--- /dev/null
+++ b/pkg/transformers/greenmask/greenmask_choice_transformer.go
@@ -0,0 +1,64 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package greenmask
+
+import (
+	"errors"
+	"fmt"
+
+	greenmasktransformers "github.com/eminano/greenmask/pkg/generators/transformers"
+	"github.com/eminano/greenmask/pkg/toolkit"
+	"github.com/xataio/pgstream/pkg/transformers"
+)
+
+type ChoiceTransformer struct {
+	transformer *greenmasktransformers.RandomChoiceTransformer
+}
+
+var errChoicesEmpty = errors.New("greenmask_choice: choices must not be empty")
+
+func NewChoiceTransformer(generatorType transformers.GeneratorType, params transformers.Parameters) (*ChoiceTransformer, error) {
+	choices := []string{}
+	choices, err := findParameter(params, "choices", choices)
+	if err != nil {
+		return nil, fmt.Errorf("greenmask_choice: choices must be an array: %w", err)
+	}
+	if len(choices) == 0 {
+		return nil, errChoicesEmpty
+	}
+
+	choicesRaw := make([]*toolkit.RawValue, len(choices))
+	for i, choice := range choices {
+		choicesRaw[i] = &toolkit.RawValue{
+			Data:   []byte(choice),
+			IsNull: false,
+		}
+	}
+
+	t := greenmasktransformers.NewRandomChoiceTransformer(choicesRaw)
+	if err := setGenerator(t, generatorType); err != nil {
+		return nil, err
+	}
+	return &ChoiceTransformer{
+		transformer: t,
+	}, nil
+}
+
+func (t *ChoiceTransformer) Transform(value any) (any, error) {
+	var toTransform []byte
+	switch val := value.(type) {
+	case []byte:
+		toTransform = val
+	case string:
+		toTransform = []byte(val)
+	case *toolkit.RawValue:
+		toTransform = val.Data
+	default:
+		return nil, transformers.ErrUnsupportedValueType
+	}
+	ret, err := t.transformer.Transform(toTransform)
+	if err != nil {
+		return nil, err
+	}
+	return ret, nil
+}
diff --git a/pkg/transformers/transformer.go b/pkg/transformers/transformer.go
index c7b57f0..0133fa9 100644
--- a/pkg/transformers/transformer.go
+++ b/pkg/transformers/transformer.go
@@ -28,6 +28,7 @@ const (
 	GreenmaskFloat     TransformerType = "greenmask_float"
 	GreenmaskUUID      TransformerType = "greenmask_uuid"
 	GreenmaskBoolean   TransformerType = "greenmask_boolean"
+	GreenmaskChoice    TransformerType = "greenmask_choice"
 )
 
 type GeneratorType string


Test patch:
diff --git a/pkg/transformers/greenmask/greenmask_choice_transformer_test.go b/pkg/transformers/greenmask/greenmask_choice_transformer_test.go
new file mode 100644
index 0000000..6409e8f
--- /dev/null
+++ b/pkg/transformers/greenmask/greenmask_choice_transformer_test.go
@@ -0,0 +1,129 @@
+// SPDX-License-Identifier: Apache-2.0
+
+package greenmask
+
+import (
+	"testing"
+
+	"github.com/eminano/greenmask/pkg/toolkit"
+	"github.com/stretchr/testify/require"
+	"github.com/xataio/pgstream/pkg/transformers"
+)
+
+func TestNewChoiceTransformer(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name      string
+		generator transformers.GeneratorType
+		params    transformers.Parameters
+		wantErr   error
+	}{
+		{
+			name:      "ok - valid random",
+			generator: transformers.Random,
+			params: transformers.Parameters{
+				"choices": []string{"a", "b", "c", "d"},
+			},
+			wantErr: nil,
+		},
+		{
+			name:      "error - invalid generator type",
+			generator: "invalid",
+			params: transformers.Parameters{
+				"choices": []string{"a", "b", "c", "d"},
+			},
+			wantErr: transformers.ErrUnsupportedGenerator,
+		},
+		{
+			name:      "error - invalid choices",
+			generator: transformers.Deterministic,
+			wantErr:   errChoicesEmpty,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			transformer, err := NewChoiceTransformer(tt.generator, tt.params)
+			require.Equal(t, tt.wantErr, err)
+			if err != nil {
+				return
+			}
+			require.NotNil(t, transformer)
+		})
+	}
+}
+
+func TestChoiceTransformer_Transform(t *testing.T) {
+	t.Parallel()
+	tests := []struct {
+		name          string
+		generatorType transformers.GeneratorType
+		input         any
+		params        transformers.Parameters
+		wantErr       error
+	}{
+		{
+			name:          "ok - transform string randomly",
+			generatorType: transformers.Random,
+			input:         "test",
+			params: transformers.Parameters{
+				"choices": []string{"a", "b", "c", "d"},
+			},
+			wantErr: nil,
+		},
+		{
+			name:          "ok - transform []byte deterministically",
+			generatorType: transformers.Deterministic,
+			input:         []byte("test"),
+			params: transformers.Parameters{
+				"choices": []string{"a", "b", "c", "d"},
+			},
+			wantErr: nil,
+		},
+		{
+			name:          "ok - transform RawValue deterministically",
+			generatorType: transformers.Deterministic,
+			input:         toolkit.NewRawValue([]byte("test"), false),
+			params: transformers.Parameters{
+				"choices": []string{"a", "b", "c", "d"},
+			},
+			wantErr: nil,
+		},
+		{
+			name:          "error - invalid input type",
+			generatorType: transformers.Random,
+			input:         1,
+			params: transformers.Parameters{
+				"choices": []string{"a", "b", "c", "d"},
+			},
+			wantErr: transformers.ErrUnsupportedValueType,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			transformer, err := NewChoiceTransformer(tt.generatorType, tt.params)
+			require.NoError(t, err)
+			require.NotNil(t, transformer)
+			got, err := transformer.Transform(tt.input)
+			require.Equal(t, tt.wantErr, err)
+			if err != nil {
+				return
+			}
+
+			// check if the result is in the choices
+			require.NotNil(t, got)
+			val, ok := got.(*toolkit.RawValue)
+			require.True(t, ok)
+			require.Contains(t, tt.params["choices"], string(val.Data))
+
+			// if deterministic, check if we get the same result again
+			if tt.generatorType == transformers.Deterministic {
+				gotAgain, err := transformer.Transform(tt.input)
+				require.NoError(t, err)
+				require.Equal(t, got, gotAgain)
+			}
+		})
+	}
+}