Repository: GreenmaskIO/greenmask
Pull request #162 (2024-08-15 09:45:24)

feat: Introduced keep_null parameter for RandomPerson
feat: Introduced keep_null parameter for columns in RandomPerson transformer


Patch:
diff --git a/docs/built_in_transformers/standard_transformers/random_person.md b/docs/built_in_transformers/standard_transformers/random_person.md
index 84c003e..1ae4220 100644
--- a/docs/built_in_transformers/standard_transformers/random_person.md
+++ b/docs/built_in_transformers/standard_transformers/random_person.md
@@ -28,6 +28,7 @@ simulating a variety of first names without using real user data.
 * `hashing` - the bool value. Indicates whether the column value must be passed through the hashing function.
   The default value is `false`. If all column has `hashing` set to `false` (by default), then all columns will be
   hashed.
+* `keep_null` - the bool value. Indicates whether NULL values should be preserved. The default value is `true`
 
 ### *gender_mapping* object attributes
 
diff --git a/internal/db/postgres/transformers/random_person.go b/internal/db/postgres/transformers/random_person.go
index fba8606..2c44c88 100644
--- a/internal/db/postgres/transformers/random_person.go
+++ b/internal/db/postgres/transformers/random_person.go
@@ -62,6 +62,7 @@ type randomNameColumns struct {
 	Name      string `json:"name"`
 	Template  string `json:"template"`
 	Hashing   bool   `json:"hashing"`
+	KeepNull  *bool  `json:"keep_null"`
 	tmpl      *template.Template
 	columnIdx int
 }
@@ -79,6 +80,7 @@ type RandomNameTransformer struct {
 	originalData []byte
 	engine       int
 	buf          *bytes.Buffer
+	nullableMap  map[int]bool
 }
 
 func NewRandomNameTransformer(ctx context.Context, driver *toolkit.Driver, parameters map[string]toolkit.Parameterizer) (utils.Transformer, toolkit.ValidationWarnings, error) {
@@ -178,6 +180,7 @@ func NewRandomNameTransformer(ctx context.Context, driver *toolkit.Driver, param
 		originalData:    make([]byte, 256),
 		engine:          engineMode,
 		buf:             bytes.NewBuffer(nil),
+		nullableMap:     make(map[int]bool, len(columns)),
 	}, nil, nil
 }
 
@@ -210,14 +213,15 @@ func (nft *RandomNameTransformer) Transform(ctx context.Context, r *toolkit.Reco
 	if nft.engine == hashEngineMode {
 		clear(nft.originalData)
 		for _, c := range nft.columns {
-			// we need to hash only columns that are marked for hashing
-			if !c.Hashing {
-				continue
-			}
 			rawVal, err := r.GetRawColumnValueByIdx(c.columnIdx)
 			if err != nil {
 				return nil, fmt.Errorf("unable to get raw value by idx %d: %w", c.columnIdx, err)
 			}
+			nft.nullableMap[c.columnIdx] = rawVal.IsNull
+			// we need to hash only columns that are marked for hashing
+			if !c.Hashing {
+				continue
+			}
 			if !rawVal.IsNull {
 				nft.originalData = append(nft.originalData, rawVal.Data...)
 			}
@@ -230,6 +234,9 @@ func (nft *RandomNameTransformer) Transform(ctx context.Context, r *toolkit.Reco
 	}
 
 	for _, c := range nft.columns {
+		if nft.nullableMap[c.columnIdx] && c.KeepNull != nil && *c.KeepNull {
+			continue
+		}
 		newRawVal := toolkit.NewRawValue(nil, false)
 		nft.buf.Reset()
 		err = c.tmpl.Execute(nft.buf, nameAttrs)
@@ -315,6 +322,11 @@ func validateColumnsAndSetDefault(driver *toolkit.Driver, columns []*randomNameC
 			c.tmpl = tmpl
 		}
 
+		if c.KeepNull == nil {
+			defaultKeepNullValue := true
+			c.KeepNull = &defaultKeepNullValue
+		}
+
 		// Do we need to calculate hash for this column?
 		if c.Hashing {
 			hasHashingColumns = true


Test patch:
diff --git a/internal/db/postgres/transformers/random_person_test.go b/internal/db/postgres/transformers/random_person_test.go
index 286d7cd..ac0b69a 100644
--- a/internal/db/postgres/transformers/random_person_test.go
+++ b/internal/db/postgres/transformers/random_person_test.go
@@ -131,3 +131,36 @@ func testStringContainsOneOfItemFromList(val string, values []string) bool {
 	}
 	return false
 }
+
+func TestRandomPersonTransformer_Transform_static_nullable(t *testing.T) {
+	columnName := "data"
+	originalValue := "\\N"
+	params := map[string]toolkit.ParamsValue{
+		"columns": toolkit.ParamsValue(`[{"name": "data", "template": "{{ .Title }} {{ .FirstName }} {{ .LastName }}"}]`),
+		"engine":  toolkit.ParamsValue("hash"),
+		"gender":  toolkit.ParamsValue("Any"),
+	}
+
+	driver, record := getDriverAndRecord(columnName, originalValue)
+	def, ok := utils.DefaultTransformerRegistry.Get("RandomPerson")
+	require.True(t, ok)
+
+	transformer, warnings, err := def.Instance(
+		context.Background(),
+		driver,
+		params,
+		nil,
+	)
+	require.NoError(t, err)
+	require.Empty(t, warnings)
+
+	r, err := transformer.Transformer.Transform(
+		context.Background(),
+		record,
+	)
+	require.NoError(t, err)
+
+	rawVal, err := r.GetRawColumnValueByName(columnName)
+	require.NoError(t, err)
+	require.True(t, rawVal.IsNull)
+}