Repository: GreenmaskIO/greenmask
Pull request #120 (2024-05-15 19:51:44)

fix: Added column size calculation based on type or column size
* Implemented GetColumnSize method for Column object
* Fixed RandomInt, NoiseInt, NoiseFloat, RandomFloat implementation
* Fixed unit tests


Patch:
diff --git a/internal/db/postgres/transformers/noise_float.go b/internal/db/postgres/transformers/noise_float.go
index adbfeae..29bc486 100644
--- a/internal/db/postgres/transformers/noise_float.go
+++ b/internal/db/postgres/transformers/noise_float.go
@@ -96,7 +96,6 @@ func NewNoiseFloatTransformer(ctx context.Context, driver *toolkit.Driver, param
 	var dynamicMode bool
 	var minValueThreshold, maxValueThreshold, minRatio, maxRatio float64
 	var precision int
-	floatSize := 8
 
 	columnParam := parameters["column"]
 	minParam := parameters["min"]
@@ -124,9 +123,7 @@ func NewNoiseFloatTransformer(ctx context.Context, driver *toolkit.Driver, param
 	}
 	affectedColumns := make(map[int]string)
 	affectedColumns[idx] = columnName
-	if c.Length != -1 {
-		floatSize = c.Length
-	}
+	floatSize := c.GetColumnSize()
 
 	if !dynamicMode {
 		if err := minParam.Scan(&minValueThreshold); err != nil {
diff --git a/internal/db/postgres/transformers/noise_int.go b/internal/db/postgres/transformers/noise_int.go
index 03d1864..f166a18 100644
--- a/internal/db/postgres/transformers/noise_int.go
+++ b/internal/db/postgres/transformers/noise_int.go
@@ -85,7 +85,6 @@ func NewNoiseIntTransformer(ctx context.Context, driver *toolkit.Driver, paramet
 	var columnName, engine string
 	var minRatio, maxRatio float64
 	var maxValueThreshold, minValueThreshold int64
-	var intSize int
 	var dynamicMode bool
 
 	columnParam := parameters["column"]
@@ -110,9 +109,7 @@ func NewNoiseIntTransformer(ctx context.Context, driver *toolkit.Driver, paramet
 	affectedColumns := make(map[int]string)
 	affectedColumns[idx] = columnName
 
-	if c.Length != -1 {
-		intSize = c.Length
-	}
+	intSize := c.GetColumnSize()
 
 	if minParam.IsDynamic() || maxParam.IsDynamic() {
 		dynamicMode = true
diff --git a/internal/db/postgres/transformers/random_float.go b/internal/db/postgres/transformers/random_float.go
index f2d1b8f..666594a 100644
--- a/internal/db/postgres/transformers/random_float.go
+++ b/internal/db/postgres/transformers/random_float.go
@@ -102,7 +102,6 @@ func NewFloatTransformer(ctx context.Context, driver *toolkit.Driver, parameters
 	var minVal, maxVal float64
 	var keepNull, dynamicMode bool
 	var precision int
-	floatSize := 8
 
 	columnParam := parameters["column"]
 	minParam := parameters["min"]
@@ -129,9 +128,7 @@ func NewFloatTransformer(ctx context.Context, driver *toolkit.Driver, parameters
 	}
 	affectedColumns := make(map[int]string)
 	affectedColumns[idx] = columnName
-	if c.Length != -1 {
-		floatSize = c.Length
-	}
+	floatSize := c.GetColumnSize()
 
 	if err := keepNullParam.Scan(&keepNull); err != nil {
 		return nil, nil, fmt.Errorf(`unable to scan "keep_null" param: %w`, err)
diff --git a/internal/db/postgres/transformers/random_int.go b/internal/db/postgres/transformers/random_int.go
index 7f0658a..b9a5bcd 100644
--- a/internal/db/postgres/transformers/random_int.go
+++ b/internal/db/postgres/transformers/random_int.go
@@ -93,7 +93,6 @@ func NewIntegerTransformer(ctx context.Context, driver *toolkit.Driver, paramete
 	var columnName, engine string
 	var minVal, maxVal int64
 	var keepNull, dynamicMode bool
-	var intSize = 8
 
 	columnParam := parameters["column"]
 	minParam := parameters["min"]
@@ -119,9 +118,7 @@ func NewIntegerTransformer(ctx context.Context, driver *toolkit.Driver, paramete
 	}
 	affectedColumns := make(map[int]string)
 	affectedColumns[idx] = columnName
-	if c.Length != -1 {
-		intSize = c.Length
-	}
+	intSize := c.GetColumnSize()
 
 	if err := keepNullParam.Scan(&keepNull); err != nil {
 		return nil, nil, fmt.Errorf(`unable to scan "keep_null" param: %w`, err)
diff --git a/pkg/toolkit/column.go b/pkg/toolkit/column.go
index ef0c91b..5030464 100644
--- a/pkg/toolkit/column.go
+++ b/pkg/toolkit/column.go
@@ -30,3 +30,10 @@ type Column struct {
 	OverriddenTypeName string `json:"overridden_type_name"`
 	OverriddenTypeOid  Oid    `json:"overridden_type_oid"`
 }
+
+func (c *Column) GetColumnSize() int {
+	if c.Length != -1 {
+		return c.Length
+	}
+	return c.TypeLength
+}


Test patch:
diff --git a/internal/db/postgres/transformers/tesing_helpers_test.go b/internal/db/postgres/transformers/tesing_helpers_test.go
index 8b88190..0408c56 100644
--- a/internal/db/postgres/transformers/tesing_helpers_test.go
+++ b/internal/db/postgres/transformers/tesing_helpers_test.go
@@ -25,84 +25,94 @@ import (
 
 var columnList = []*toolkit.Column{
 	{
-		Name:     "id",
-		TypeName: "int2",
-		TypeOid:  pgtype.Int2OID,
-		Num:      1,
-		NotNull:  false,
-		Length:   2,
-	},
-	{
-		Name:     "created_at",
-		TypeName: "timestamp",
-		TypeOid:  pgtype.TimestampOID,
-		Num:      2,
-		NotNull:  false,
-		Length:   -1,
-	},
-	{
-		Name:     "data",
-		TypeName: "text",
-		TypeOid:  pgtype.TextOID,
-		Num:      3,
-		NotNull:  false,
-		Length:   -1,
-	},
-	{
-		Name:     "doc",
-		TypeName: "jsonb",
-		TypeOid:  pgtype.JSONBOID,
-		Num:      4,
-		NotNull:  false,
-		Length:   -1,
-	},
-	{
-		Name:     "uid",
-		TypeName: "uuid",
-		TypeOid:  pgtype.UUIDOID,
-		Num:      5,
-		NotNull:  false,
-		Length:   -1,
-	},
-	{
-		Name:     "id2",
-		TypeName: "int2",
-		TypeOid:  pgtype.Int2OID,
-		Num:      6,
-		NotNull:  false,
-		Length:   2,
-	},
-	{
-		Name:     "id4",
-		TypeName: "int4",
-		TypeOid:  pgtype.Int4OID,
-		Num:      7,
-		NotNull:  false,
-		Length:   4,
-	},
-	{
-		Name:     "id8",
-		TypeName: "int8",
-		TypeOid:  pgtype.Int8OID,
-		Num:      8,
-		NotNull:  false,
-		Length:   8,
-	},
-	{
-		Name:     "date_date",
-		TypeName: "date",
-		TypeOid:  pgtype.DateOID,
-		Num:      9,
-		NotNull:  false,
-		Length:   -1,
-	},
-	{
-		Name:     "date_ts",
-		TypeName: "timestamp",
-		TypeOid:  pgtype.TimestampOID,
-		Num:      10,
-		NotNull:  false,
-		Length:   -1,
+		Name:       "id",
+		TypeName:   "int2",
+		TypeOid:    pgtype.Int2OID,
+		Num:        1,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: 2,
+	},
+	{
+		Name:       "created_at",
+		TypeName:   "timestamp",
+		TypeOid:    pgtype.TimestampOID,
+		Num:        2,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: -1,
+	},
+	{
+		Name:       "data",
+		TypeName:   "text",
+		TypeOid:    pgtype.TextOID,
+		Num:        3,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: -1,
+	},
+	{
+		Name:       "doc",
+		TypeName:   "jsonb",
+		TypeOid:    pgtype.JSONBOID,
+		Num:        4,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: -1,
+	},
+	{
+		Name:       "uid",
+		TypeName:   "uuid",
+		TypeOid:    pgtype.UUIDOID,
+		Num:        5,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: -1,
+	},
+	{
+		Name:       "id2",
+		TypeName:   "int2",
+		TypeOid:    pgtype.Int2OID,
+		Num:        6,
+		NotNull:    false,
+		TypeLength: 2,
+		Length:     -1,
+	},
+	{
+		Name:       "id4",
+		TypeName:   "int4",
+		TypeOid:    pgtype.Int4OID,
+		Num:        7,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: 4,
+	},
+	{
+		Name:       "id8",
+		TypeName:   "int8",
+		TypeOid:    pgtype.Int8OID,
+		Num:        8,
+		NotNull:    false,
+		TypeLength: 8,
+		Length:     -1,
+	},
+	{
+		Name:       "date_date",
+		TypeName:   "date",
+		TypeOid:    pgtype.DateOID,
+		Num:        9,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: -1,
+	},
+	{
+		Name:       "date_ts",
+		TypeName:   "timestamp",
+		TypeOid:    pgtype.TimestampOID,
+		Num:        10,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: -1,
 	},
 	{
 		Name:     "date_tstz",
@@ -113,20 +123,22 @@ var columnList = []*toolkit.Column{
 		Length:   -1,
 	},
 	{
-		Name:     "col_float4",
-		TypeName: "float4",
-		TypeOid:  pgtype.Float4OID,
-		Num:      12,
-		NotNull:  false,
-		Length:   -1,
+		Name:       "col_float4",
+		TypeName:   "float4",
+		TypeOid:    pgtype.Float4OID,
+		Num:        12,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: 4,
 	},
 	{
-		Name:     "col_float8",
-		TypeName: "float8",
-		TypeOid:  pgtype.Float8OID,
-		Num:      13,
-		NotNull:  false,
-		Length:   -1,
+		Name:       "col_float8",
+		TypeName:   "float8",
+		TypeOid:    pgtype.Float8OID,
+		Num:        13,
+		NotNull:    false,
+		Length:     -1,
+		TypeLength: 8,
 	},
 	{
 		Name:     "col_bool",