Merge pull request #119 from GreenmaskIO/fix/noise_numeric_limiter_issue

fix: Fixed NoiseNumeric and RandomNumeric limiters

internal/db/postgres/transformers/noise_numeric.go

@@ -107,7 +107,7 @@ type NoiseNumericTransformer struct {
 	maxRatioParam  toolkit.Parameterizer
 	minRatioParam  toolkit.Parameterizer
 
-	transform func(context.Context, decimal.Decimal) (decimal.Decimal, error)
+	transform func(decimal.Decimal) (decimal.Decimal, error)
 }
 
 func NewNumericFloatTransformer(ctx context.Context, driver *toolkit.Driver, parameters map[string]toolkit.Parameterizer) (utils.Transformer, toolkit.ValidationWarnings, error) {
@@ -224,7 +224,7 @@ func (nft *NoiseNumericTransformer) Done(ctx context.Context) error {
 	return nil
 }
 
-func (nft *NoiseNumericTransformer) dynamicTransform(ctx context.Context, original decimal.Decimal) (decimal.Decimal, error) {
+func (nft *NoiseNumericTransformer) dynamicTransform(original decimal.Decimal) (decimal.Decimal, error) {
 	var minVal, maxVal decimal.Decimal
 	err := nft.minParam.Scan(&minVal)
 	if err != nil {
@@ -236,12 +236,12 @@ func (nft *NoiseNumericTransformer) dynamicTransform(ctx context.Context, origin
 		return decimal.Decimal{}, fmt.Errorf(`unable to scan "max" param: %w`, err)
 	}
 
-	limiter, err := getNumericLimiterForDynamicParameter(nft.numericSize, minVal, maxVal, nft.minAllowedValue, nft.maxAllowedValue)
+	limiter, err := getNoiseNumericLimiterForDynamicParameter(nft.numericSize, minVal, maxVal, nft.minAllowedValue, nft.maxAllowedValue)
 	if err != nil {
 		return decimal.Decimal{}, fmt.Errorf("error creating limiter in dynamic mode: %w", err)
 	}
-	ctx = context.WithValue(ctx, "limiter", limiter)
-	return nft.t.Transform(ctx, original)
+	limiter.SetPrecision(nft.precision)
+	return nft.t.SetDynamicLimiter(limiter).Transform(original)
 }
 
 func (nft *NoiseNumericTransformer) Transform(ctx context.Context, r *toolkit.Record) (*toolkit.Record, error) {
@@ -254,7 +254,7 @@ func (nft *NoiseNumericTransformer) Transform(ctx context.Context, r *toolkit.Re
 		return r, nil
 	}
 
-	res, err := nft.transform(ctx, val)
+	res, err := nft.transform(val)
 	if err != nil {
 		return nil, fmt.Errorf("unable to transform value: %w", err)
 	}
@@ -285,6 +285,33 @@ func validateNoiseNumericTypeAndSetLimit(
 	return limiter, nil, nil
 }
 
+func getNoiseNumericLimiterForDynamicParameter(
+	numericSize int, requestedMinValue, requestedMaxValue,
+	minAllowedValue, maxAllowedValue decimal.Decimal,
+) (*transformers.NoiseNumericLimiter, error) {
+
+	if !numericLimitIsValid(requestedMinValue, minAllowedValue, maxAllowedValue) {
+		return nil, fmt.Errorf("requested dynamic parameter min value is out of range of NUMERIC(%d) size", numericSize)
+	}
+
+	if !numericLimitIsValid(requestedMaxValue, minAllowedValue, maxAllowedValue) {
+		return nil, fmt.Errorf("requested dynamic parameter max value is out of range of NUMERIC(%d) size", numericSize)
+	}
+
+	limiter, err := transformers.NewNoiseNumericLimiter(minAllowedValue, maxAllowedValue)
+	if err != nil {
+		return nil, err
+	}
+
+	if !requestedMinValue.Equal(decimal.NewFromInt(0)) || !requestedMinValue.Equal(decimal.NewFromInt(0)) {
+		limiter, err = transformers.NewNoiseNumericLimiter(requestedMinValue, requestedMaxValue)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return limiter, nil
+}
+
 func init() {
 	utils.DefaultTransformerRegistry.MustRegister(NoiseNumericTransformerDefinition)
 }

internal/db/postgres/transformers/random_numeric.go

@@ -97,7 +97,7 @@ type NumericTransformer struct {
 	keepNullParam  toolkit.Parameterizer
 	engineParam    toolkit.Parameterizer
 	precisionParam toolkit.Parameterizer
-	transform      func(context.Context, []byte) (decimal.Decimal, error)
+	transform      func([]byte) (decimal.Decimal, error)
 }
 
 func NewRandomNumericTransformer(ctx context.Context, driver *toolkit.Driver, parameters map[string]toolkit.Parameterizer) (utils.Transformer, toolkit.ValidationWarnings, error) {
@@ -208,7 +208,7 @@ func (bit *NumericTransformer) Done(ctx context.Context) error {
 	return nil
 }
 
-func (bit *NumericTransformer) dynamicTransform(ctx context.Context, v []byte) (decimal.Decimal, error) {
+func (bit *NumericTransformer) dynamicTransform(v []byte) (decimal.Decimal, error) {
 	var minVal, maxVal decimal.Decimal
 	err := bit.minParam.Scan(&minVal)
 	if err != nil {
@@ -220,12 +220,11 @@ func (bit *NumericTransformer) dynamicTransform(ctx context.Context, v []byte) (
 		return decimal.Decimal{}, fmt.Errorf(`unable to scan "max" param: %w`, err)
 	}
 
-	limiter, err := getNumericLimiterForDynamicParameter(bit.numericSize, minVal, maxVal, bit.minAllowedValue, bit.maxAllowedValue)
+	limiter, err := getRandomNumericLimiterForDynamicParameter(bit.numericSize, minVal, maxVal, bit.minAllowedValue, bit.maxAllowedValue)
 	if err != nil {
 		return decimal.Decimal{}, fmt.Errorf("error creating limiter in dynamic mode: %w", err)
 	}
-	ctx = context.WithValue(ctx, "limiter", limiter)
-	return bit.RandomNumericTransformer.Transform(ctx, v)
+	return bit.RandomNumericTransformer.SetDynamicLimiter(limiter).Transform(v)
 }
 
 func (bit *NumericTransformer) Transform(ctx context.Context, r *toolkit.Record) (*toolkit.Record, error) {
@@ -237,7 +236,7 @@ func (bit *NumericTransformer) Transform(ctx context.Context, r *toolkit.Record)
 		return r, nil
 	}
 
-	newValue, err := bit.transform(ctx, val.Data)
+	newValue, err := bit.transform(val.Data)
 	if err != nil {
 		return nil, err
 	}
@@ -311,7 +310,7 @@ func numericLimitIsValid(requestedThreshold, minValue, maxValue decimal.Decimal)
 	return requestedThreshold.GreaterThanOrEqual(minValue) || requestedThreshold.LessThanOrEqual(maxValue)
 }
 
-func getNumericLimiterForDynamicParameter(
+func getRandomNumericLimiterForDynamicParameter(
 	numericSize int, requestedMinValue, requestedMaxValue,
 	minAllowedValue, maxAllowedValue decimal.Decimal,
 ) (*transformers.RandomNumericLimiter, error) {

internal/generators/transformers/noise_numeric.go

@@ -1,7 +1,6 @@
 package transformers
 
 import (
-	"context"
 	"fmt"
 
 	"github.com/greenmaskio/greenmask/internal/generators"
@@ -30,6 +29,12 @@ func NewNoiseNumericLimiter(minVal, maxVal decimal.Decimal) (*NoiseNumericLimite
 	}, nil
 }
 
+func (ni *NoiseNumericLimiter) SetPrecision(v int32) *NoiseNumericLimiter {
+	ni.precision = v
+	ni.withPrecision = true
+	return ni
+}
+
 func (ni *NoiseNumericLimiter) Limit(v decimal.Decimal) decimal.Decimal {
 	if v.GreaterThan(ni.MaxValue) {
 		return ni.MaxValue
@@ -43,17 +48,13 @@ func (ni *NoiseNumericLimiter) Limit(v decimal.Decimal) decimal.Decimal {
 	return v
 }
 
-func (ni *NoiseNumericLimiter) SetPrecision(v int32) {
-	ni.precision = v
-	ni.withPrecision = true
-}
-
 type NoiseNumericTransformer struct {
-	generator  generators.Generator
-	limiter    *NoiseNumericLimiter
-	byteLength int
-	minRatio   float64
-	maxRatio   float64
+	generator      generators.Generator
+	limiter        *NoiseNumericLimiter
+	dynamicLimiter *NoiseNumericLimiter
+	byteLength     int
+	minRatio       float64
+	maxRatio       float64
 }
 
 func NewNoiseNumericTransformer(limiter *NoiseNumericLimiter, minRatio, maxRatio float64) *NoiseNumericTransformer {
@@ -65,12 +66,21 @@ func NewNoiseNumericTransformer(limiter *NoiseNumericLimiter, minRatio, maxRatio
 	}
 }
 
-func (nt *NoiseNumericTransformer) Transform(ctx context.Context, original decimal.Decimal) (decimal.Decimal, error) {
-	var limiter = nt.limiter
-	limiterAny := ctx.Value("limiter")
+// SetDynamicLimiter sets the limiter for the dynamic mode. dynamicLimiter will be used set as nil after the Transform
+// call.
+func (nt *NoiseNumericTransformer) SetDynamicLimiter(l *NoiseNumericLimiter) *NoiseNumericTransformer {
+	if l == nil {
+		panic("bug: limiter for NoiseNumericTransformer values cannot be nil")
+	}
+	nt.dynamicLimiter = l
+	return nt
+}
+
+func (nt *NoiseNumericTransformer) Transform(original decimal.Decimal) (decimal.Decimal, error) {
 
-	if limiterAny != nil {
-		limiter = limiterAny.(*NoiseNumericLimiter)
+	limiter := nt.limiter
+	if nt.dynamicLimiter != nil {
+		limiter = nt.dynamicLimiter
 	}
 
 	resBytes, err := nt.generator.Generate(original.BigInt().Bytes())
@@ -96,6 +106,10 @@ func (nt *NoiseNumericTransformer) Transform(ctx context.Context, original decim
 		res = limiter.Limit(res)
 	}
 
+	if nt.dynamicLimiter != nil {
+		limiter = nil
+	}
+
 	return res, nil
 }
 

internal/generators/transformers/random_numeric.go

@@ -1,7 +1,6 @@
 package transformers
 
 import (
-	"context"
 	"fmt"
 	"math/big"
 	"strings"
@@ -82,10 +81,11 @@ func (l *RandomNumericLimiter) Limit(v decimal.Decimal) decimal.Decimal {
 }
 
 type RandomNumericTransformer struct {
-	generator  generators.Generator
-	limiter    *RandomNumericLimiter
-	byteLength int
-	precision  int32
+	generator      generators.Generator
+	limiter        *RandomNumericLimiter
+	dynamicLimiter *RandomNumericLimiter
+	byteLength     int
+	precision      int32
 }
 
 func NewRandomNumericTransformer(limiter *RandomNumericLimiter, precision int32) (*RandomNumericTransformer, error) {
@@ -114,13 +114,21 @@ func (ig *RandomNumericTransformer) SetGenerator(g generators.Generator) error {
 	return nil
 }
 
-func (ig *RandomNumericTransformer) Transform(ctx context.Context, original []byte) (decimal.Decimal, error) {
+// SetDynamicLimiter sets the limiter for the dynamic mode. dynamicLimiter will be used set as nil after the Transform
+// call.
+func (ig *RandomNumericTransformer) SetDynamicLimiter(l *RandomNumericLimiter) *RandomNumericTransformer {
+	if l == nil {
+		panic("bug: limiter for RandomNumericTransformer values cannot be nil")
+	}
+	ig.dynamicLimiter = l
+	return ig
+}
+
+func (ig *RandomNumericTransformer) Transform(original []byte) (decimal.Decimal, error) {
 	var res decimal.Decimal
 	var limiter = ig.limiter
-	limiterAny := ctx.Value("limiter")
-
-	if limiterAny != nil {
-		limiter = limiterAny.(*RandomNumericLimiter)
+	if ig.dynamicLimiter != nil {
+		limiter = ig.dynamicLimiter
 	}
 
 	resBytes, err := ig.generator.Generate(original)
@@ -134,6 +142,10 @@ func (ig *RandomNumericTransformer) Transform(ctx context.Context, original []by
 	}
 	res = limiter.Limit(res)
 
+	if ig.dynamicLimiter != nil {
+		limiter = nil
+	}
+
 	return res, nil
 }
 

internal/generators/transformers/random_numeric_test.go

@@ -1,7 +1,6 @@
 package transformers
 
 import (
-	"context"
 	"testing"
 
 	"github.com/shopspring/decimal"
@@ -22,7 +21,7 @@ func TestBigIntTransformer_Transform(t *testing.T) {
 	require.NoError(t, err)
 	err = tr.SetGenerator(sha1)
 	require.NoError(t, err)
-	res, err := tr.Transform(context.Background(), []byte("199999999999999999999999999999999999999"))
+	res, err := tr.Transform([]byte("199999999999999999999999999999999999999"))
 	require.NoError(t, err)
 	require.True(t, res.LessThanOrEqual(maxValue) && res.GreaterThanOrEqual(minValue))
 }
@@ -40,7 +39,7 @@ func TestBigFloatTransformer_Transform(t *testing.T) {
 	require.NoError(t, err)
 	err = tr.SetGenerator(sha1)
 	require.NoError(t, err)
-	res, err := tr.Transform(context.Background(), []byte("1999999999999999999999999999999999999990"))
+	res, err := tr.Transform([]byte("1999999999999999999999999999999999999990"))
 	require.NoError(t, err)
 	require.True(t, res.LessThanOrEqual(maxValue) && res.GreaterThanOrEqual(minValue))
 }