Skip to content

Commit

Permalink
Clean up.
Browse files Browse the repository at this point in the history
  • Loading branch information
@Tang8330
Tang8330 committed Sep 26, 2024
1 parent e45caab commit e34defc
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 56 deletions.
134 changes: 83 additions & 51 deletions clients/redshift/cast_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,40 +3,46 @@ package redshift
import (
"fmt"

"github.com/artie-labs/transfer/lib/stringutil"

"github.com/artie-labs/transfer/lib/config/constants"

"github.com/artie-labs/transfer/lib/stringutil"
"github.com/artie-labs/transfer/lib/typing"
"github.com/stretchr/testify/assert"
)

func (r *RedshiftTestSuite) TestCanIncreasePrecision() {
{
// Not a string
assert.False(r.T(), canIncreasePrecision(typing.Struct))
}
{
// String, but precision is not specified
assert.False(r.T(), canIncreasePrecision(typing.String))
}
{
// String, but maxed out already
assert.False(r.T(), canIncreasePrecision(
typing.KindDetails{
Kind: typing.String.Kind,
OptionalStringPrecision: typing.ToPtr(maxRedshiftLength),
}),
)
// False
{
// Not a string
assert.False(r.T(), canIncreasePrecision(typing.Struct, 123))
}
{
// String, but precision not specified.
assert.False(r.T(), canIncreasePrecision(typing.String, 123))
}
{
// String and precision specified, but value length exceeds maxRedshiftLength
assert.False(r.T(), canIncreasePrecision(
typing.KindDetails{
Kind: typing.String.Kind,
OptionalStringPrecision: typing.ToPtr(int32(123)),
},
maxRedshiftLength+1),
)
}
}
{
// String, precision is low and can be increased
assert.True(r.T(), canIncreasePrecision(
typing.KindDetails{
Kind: typing.String.Kind,
OptionalStringPrecision: typing.ToPtr(maxRedshiftLength - 1),
}),
)
// True
{
// String, precision is low and can be increased
assert.True(r.T(), canIncreasePrecision(
typing.KindDetails{
Kind: typing.String.Kind,
OptionalStringPrecision: typing.ToPtr(int32(123)),
},
123),
)
}
}
}

Expand All @@ -45,58 +51,76 @@ func (r *RedshiftTestSuite) TestReplaceExceededValues() {
// Irrelevant data type
{
// Integer
assert.Equal(r.T(), "123", replaceExceededValues("123", typing.Integer, false, false))

value, _ := replaceExceededValues("123", typing.Integer, false, false)
assert.Equal(r.T(), "123", value)
}
{
// Returns the full value since it's not a struct or string
// This is invalid and should not happen, but it's here to ensure we're only checking for structs and strings.
value := stringutil.Random(int(maxRedshiftLength + 1))
assert.Equal(r.T(), value, replaceExceededValues(value, typing.Integer, false, false))
input := stringutil.Random(int(maxRedshiftLength + 1))
value, _ := replaceExceededValues(input, typing.Integer, false, false)
assert.Equal(r.T(), input, value)
}
}
{
// Exceeded
{
// String
{
// TruncateExceededValue = false, IncreaseStringPrecision = false
assert.Equal(r.T(), constants.ExceededValueMarker, replaceExceededValues(stringutil.Random(int(maxRedshiftLength)+1), typing.String, false, false))
value, shouldIncrease := replaceExceededValues(stringutil.Random(int(maxRedshiftLength)+1), typing.String, false, false)
assert.Equal(r.T(), constants.ExceededValueMarker, value)
assert.False(r.T(), shouldIncrease)
}
{
// TruncateExceededValue = false, string precision specified, IncreaseStringPrecision = false
stringKd := typing.KindDetails{
Kind: typing.String.Kind,
OptionalStringPrecision: typing.ToPtr(int32(3)),
}

assert.Equal(r.T(), constants.ExceededValueMarker, replaceExceededValues("hello", stringKd, false, false))
value, shouldIncrease := replaceExceededValues("hello", stringKd, false, false)
assert.Equal(r.T(), constants.ExceededValueMarker, value)
assert.False(r.T(), shouldIncrease)
}
{
// TruncateExceededValue = true, IncreaseStringPrecision = false
superLongString := stringutil.Random(int(maxRedshiftLength) + 1)
assert.Equal(r.T(), superLongString[:maxRedshiftLength], replaceExceededValues(superLongString, typing.String, true, false))
input := stringutil.Random(int(maxRedshiftLength) + 1)
value, shouldIncrease := replaceExceededValues(input, typing.String, true, false)
assert.Equal(r.T(), input[:maxRedshiftLength], value)
assert.False(r.T(), shouldIncrease)
}
{
// TruncateExceededValue = true, string precision specified, IncreaseStringPrecision = false
stringKd := typing.KindDetails{
Kind: typing.String.Kind,
OptionalStringPrecision: typing.ToPtr(int32(3)),
}

assert.Equal(r.T(), "hel", replaceExceededValues("hello", stringKd, true, false))
value, shouldIncrease := replaceExceededValues("hello", stringKd, true, false)
assert.Equal(r.T(), "hel", value)
assert.False(r.T(), shouldIncrease)
}
}
{
// Struct and masked
assert.Equal(r.T(), fmt.Sprintf(`{"key":"%s"}`, constants.ExceededValueMarker), replaceExceededValues(fmt.Sprintf(`{"foo": "%s"}`, stringutil.Random(int(maxRedshiftLength)+1)), typing.Struct, false, false))
value, shouldIncrease := replaceExceededValues(fmt.Sprintf(`{"foo": "%s"}`, stringutil.Random(int(maxRedshiftLength)+1)), typing.Struct, false, false)
assert.Equal(r.T(), fmt.Sprintf(`{"key":"%s"}`, constants.ExceededValueMarker), value)
assert.False(r.T(), shouldIncrease)
}
}
{
// Valid
{
// Not masked
assert.Equal(r.T(), `{"foo": "bar"}`, replaceExceededValues(`{"foo": "bar"}`, typing.Struct, false, false))
assert.Equal(r.T(), "hello world", replaceExceededValues("hello world", typing.String, false, false))
{
value, shouldIncrease := replaceExceededValues(`{"foo": "bar"}`, typing.Struct, false, false)
assert.Equal(r.T(), `{"foo": "bar"}`, value)
assert.False(r.T(), shouldIncrease)
}
{
value, shouldIncrease := replaceExceededValues("hello world", typing.String, false, false)
assert.Equal(r.T(), "hello world", value)
assert.False(r.T(), shouldIncrease)
}
}
}
}
Expand All @@ -108,16 +132,18 @@ func (r *RedshiftTestSuite) TestCastColValStaging() {
// String
{
// TruncateExceededValue = false, IncreaseStringPrecision = false
value, err := castColValStaging(stringutil.Random(int(maxRedshiftLength)+1), typing.String, false, false)
value, shouldIncrease, err := castColValStaging(stringutil.Random(int(maxRedshiftLength)+1), typing.String, false, false)
assert.NoError(r.T(), err)
assert.Equal(r.T(), constants.ExceededValueMarker, value)
assert.False(r.T(), shouldIncrease)
}
{
// TruncateExceededValue = true, IncreaseStringPrecision = false
value := stringutil.Random(int(maxRedshiftLength) + 1)
value, err := castColValStaging(value, typing.String, true, false)
input := stringutil.Random(int(maxRedshiftLength) + 1)
value, shouldIncrease, err := castColValStaging(input, typing.String, true, false)
assert.NoError(r.T(), err)
assert.Equal(r.T(), value[:maxRedshiftLength], value)
assert.Equal(r.T(), input[:maxRedshiftLength], value)
assert.False(r.T(), shouldIncrease)
}
{
// TruncateExceededValue = false, IncreaseStringPrecision = true
Expand All @@ -126,43 +152,49 @@ func (r *RedshiftTestSuite) TestCastColValStaging() {
OptionalStringPrecision: typing.ToPtr(int32(3)),
}

value, err := castColValStaging("hello", stringKd, false, true)
value, shouldIncrease, err := castColValStaging("hello", stringKd, false, true)
assert.NoError(r.T(), err)
assert.Equal(r.T(), "hello", value)
assert.True(r.T(), shouldIncrease)
}
{
value := stringutil.Random(int(maxRedshiftLength) + 1)
// TruncateExceededValue = true, IncreaseStringPrecision = true
input := stringutil.Random(int(maxRedshiftLength) + 1)
stringPrecision := int32(3)
stringKd := typing.KindDetails{
Kind: typing.String.Kind,
OptionalStringPrecision: typing.ToPtr(int32(3)),
OptionalStringPrecision: typing.ToPtr(stringPrecision),
}

value, err := castColValStaging(value, stringKd, true, true)
value, shouldIncrease, err := castColValStaging(input, stringKd, true, true)
assert.NoError(r.T(), err)
assert.Equal(r.T(), value[:maxRedshiftLength], value)
assert.Equal(r.T(), input[:stringPrecision], value)
assert.False(r.T(), shouldIncrease)
}
}
{
// Masked struct
value, err := castColValStaging(fmt.Sprintf(`{"foo": "%s"}`, stringutil.Random(int(maxRedshiftLength)+1)), typing.Struct, false, false)
value, shouldIncrease, err := castColValStaging(fmt.Sprintf(`{"foo": "%s"}`, stringutil.Random(int(maxRedshiftLength)+1)), typing.Struct, false, false)
assert.NoError(r.T(), err)
assert.Equal(r.T(), fmt.Sprintf(`{"key":"%s"}`, constants.ExceededValueMarker), value)
assert.False(r.T(), shouldIncrease)
}
}
{
// Not exceeded
{
// Valid string
value, err := castColValStaging("thisissuperlongbutnotlongenoughtogetmasked", typing.String, false, false)
value, shouldIncrease, err := castColValStaging("thisissuperlongbutnotlongenoughtogetmasked", typing.String, false, false)
assert.NoError(r.T(), err)
assert.Equal(r.T(), "thisissuperlongbutnotlongenoughtogetmasked", value)
assert.False(r.T(), shouldIncrease)
}
{
// Valid struct
value, err := castColValStaging(`{"foo": "bar"}`, typing.Struct, false, false)
value, shouldIncrease, err := castColValStaging(`{"foo": "bar"}`, typing.Struct, false, false)
assert.NoError(r.T(), err)
assert.Equal(r.T(), `{"foo": "bar"}`, value)
assert.False(r.T(), shouldIncrease)
}
}
}
7 changes: 2 additions & 5 deletions clients/redshift/staging.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,8 @@ func (s *Store) loadTemporaryTable(tableData *optimization.TableData, newTableID
for _, value := range tableData.Rows() {
var row []string
for _, col := range columns {
castedValue, shouldIncreaseColumn, castErr := castColValStaging(
// TODO: Implement
castedValue, _, castErr := castColValStaging(
value[col.Name()],
col.KindDetails,
s.config.SharedDestinationSettings.TruncateExceededValues,
Expand All @@ -104,10 +105,6 @@ func (s *Store) loadTemporaryTable(tableData *optimization.TableData, newTableID
return "", castErr
}

if shouldIncreaseColumn {
s.dialect().BuildIncreaseStringPrecisionQuery()
}

row = append(row, castedValue)
}

Expand Down

0 comments on commit e34defc

Please sign in to comment.