[Typing] Support different Redshift Integers (#926)

artie-labs · Sep 25, 2024 · 756a7fe · 756a7fe
1 parent 067f4c8
commit 756a7fe
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 15 deletions.
diff --git a/clients/redshift/dialect/dialect.go b/clients/redshift/dialect/dialect.go
@@ -26,8 +26,20 @@ func (RedshiftDialect) EscapeStruct(value string) string {
 func (RedshiftDialect) DataTypeForKind(kd typing.KindDetails, _ bool) string {
 	switch kd.Kind {
 	case typing.Integer.Kind:
-		// int4 is 2^31, whereas int8 is 2^63.
-		// we're using a larger data type to not have an integer overflow.
+		if kd.OptionalIntegerKind != nil {
+			switch *kd.OptionalIntegerKind {
+			case typing.SmallIntegerKind:
+				return "INT2"
+			case typing.IntegerKind:
+				return "INT4"
+			case typing.NotSpecifiedKind, typing.BigIntegerKind:
+				fallthrough
+			default:
+				// By default, we are using a larger data type to avoid the possibility of an integer overflow.
+				return "INT8"
+			}
+		}
+
 		return "INT8"
 	case typing.Struct.Kind:
 		return "SUPER"
@@ -87,8 +99,21 @@ func (RedshiftDialect) KindForDataType(rawType string, stringPrecision string) (
 	switch rawType {
 	case "super":
 		return typing.Struct, nil
-	case "smallint", "integer", "bigint":
-		return typing.Integer, nil
+	case "smallint":
+		return typing.KindDetails{
+			Kind:                typing.Integer.Kind,
+			OptionalIntegerKind: typing.ToPtr(typing.SmallIntegerKind),
+		}, nil
+	case "integer":
+		return typing.KindDetails{
+			Kind:                typing.Integer.Kind,
+			OptionalIntegerKind: typing.ToPtr(typing.IntegerKind),
+		}, nil
+	case "bigint":
+		return typing.KindDetails{
+			Kind:                typing.Integer.Kind,
+			OptionalIntegerKind: typing.ToPtr(typing.BigIntegerKind),
+		}, nil
 	case "double precision":
 		return typing.Float, nil
 	case "timestamp with time zone", "timestamp without time zone":
@@ -101,7 +126,7 @@ func (RedshiftDialect) KindForDataType(rawType string, stringPrecision string) (
 		return typing.Boolean, nil
 	}
 
-	return typing.Invalid, fmt.Errorf("unsupported data type: %s", rawType)
+	return typing.Invalid, fmt.Errorf("unsupported data type: %q", rawType)
 }
 
 func (RedshiftDialect) IsColumnAlreadyExistsErr(err error) bool {

diff --git a/clients/redshift/dialect/dialect_test.go b/clients/redshift/dialect/dialect_test.go
@@ -29,26 +29,62 @@ func TestRedshiftDialect_DataTypeForKind(t *testing.T) {
 			assert.Equal(t, "VARCHAR(12345)", RedshiftDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.String.Kind, OptionalStringPrecision: typing.ToPtr(int32(12345))}, false))
 		}
 	}
+	{
+		// Integers
+		{
+			// Small int
+			assert.Equal(t, "INT2", RedshiftDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.SmallIntegerKind)}, false))
+		}
+		{
+			// Integer
+			assert.Equal(t, "INT4", RedshiftDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.IntegerKind)}, false))
+		}
+		{
+			// Big integer
+			assert.Equal(t, "INT8", RedshiftDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.BigIntegerKind)}, false))
+		}
+		{
+			// Not specified
+			{
+				// Literal
+				assert.Equal(t, "INT8", RedshiftDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.NotSpecifiedKind)}, false))
+			}
+			{
+				assert.Equal(t, "INT8", RedshiftDialect{}.DataTypeForKind(typing.Integer, false))
+			}
+		}
+	}
 }
 
 func TestRedshiftDialect_KindForDataType(t *testing.T) {
 	dialect := RedshiftDialect{}
 	{
 		// Integers
 		{
-			kd, err := dialect.KindForDataType("integer", "")
+			// Small integer
+			kd, err := dialect.KindForDataType("smallint", "")
 			assert.NoError(t, err)
-			assert.Equal(t, typing.Integer, kd)
+			assert.Equal(t, typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.SmallIntegerKind)}, kd)
 		}
 		{
-			kd, err := dialect.KindForDataType("bigint", "")
-			assert.NoError(t, err)
-			assert.Equal(t, typing.Integer, kd)
+			{
+				// Regular integers (upper)
+				kd, err := dialect.KindForDataType("INTEGER", "")
+				assert.NoError(t, err)
+				assert.Equal(t, typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.IntegerKind)}, kd)
+			}
+			{
+				// Regular integers (lower)
+				kd, err := dialect.KindForDataType("integer", "")
+				assert.NoError(t, err)
+				assert.Equal(t, typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.IntegerKind)}, kd)
+			}
 		}
 		{
-			kd, err := dialect.KindForDataType("INTEGER", "")
+			// Big integer
+			kd, err := dialect.KindForDataType("bigint", "")
 			assert.NoError(t, err)
-			assert.Equal(t, typing.Integer, kd)
+			assert.Equal(t, typing.KindDetails{Kind: typing.Integer.Kind, OptionalIntegerKind: typing.ToPtr(typing.BigIntegerKind)}, kd)
 		}
 	}
 	{

diff --git a/lib/optimization/table_data.go b/lib/optimization/table_data.go
@@ -282,6 +282,11 @@ func (t *TableData) MergeColumnsFromDestination(destCols ...columns.Column) erro
 				inMemoryCol.KindDetails.OptionalStringPrecision = foundColumn.KindDetails.OptionalStringPrecision
 			}
 
+			// Copy over integer kind, if exists.
+			if foundColumn.KindDetails.OptionalIntegerKind != nil {
+				inMemoryCol.KindDetails.OptionalIntegerKind = foundColumn.KindDetails.OptionalIntegerKind
+			}
+
 			// Copy over the time details
 			if foundColumn.KindDetails.ExtendedTimeDetails != nil {
 				if inMemoryCol.KindDetails.ExtendedTimeDetails == nil {

diff --git a/lib/typing/typing.go b/lib/typing/typing.go
@@ -9,13 +9,23 @@ import (
 	"github.com/artie-labs/transfer/lib/typing/ext"
 )
 
+type OptionalIntegerKind int
+
+const (
+	NotSpecifiedKind OptionalIntegerKind = iota
+	SmallIntegerKind
+	IntegerKind
+	BigIntegerKind
+)
+
 type KindDetails struct {
 	Kind                   string
 	ExtendedTimeDetails    *ext.NestedKind
 	ExtendedDecimalDetails *decimal.Details
 
 	// Optional kind details metadata
 	OptionalStringPrecision *int32
+	OptionalIntegerKind     *OptionalIntegerKind
 }
 
 func (k *KindDetails) EnsureExtendedTimeDetails() error {
@@ -26,8 +36,6 @@ func (k *KindDetails) EnsureExtendedTimeDetails() error {
 	return nil
 }
 
-// Summarized this from Snowflake + Reflect.
-// In the future, we can support Geo objects.
 var (
 	Invalid = KindDetails{
 		Kind: "invalid",
@@ -38,7 +46,8 @@ var (
 	}
 
 	Integer = KindDetails{
-		Kind: "int",
+		Kind:                "int",
+		OptionalIntegerKind: ToPtr(NotSpecifiedKind),
 	}
 
 	EDecimal = KindDetails{