WIP.

artie-labs · Sep 30, 2024 · 6d3cc45 · 6d3cc45
1 parent 339c0e5
commit 6d3cc45
Show file tree

Hide file tree

Showing 4 changed files with 79 additions and 2 deletions.
diff --git a/clients/databricks/dialect/typing.go b/clients/databricks/dialect/typing.go
@@ -9,7 +9,7 @@ import (
 	"github.com/artie-labs/transfer/lib/typing/ext"
 )
 
-func (DatabricksDialect) DataTypeForKind(kindDetails typing.KindDetails, isPk bool) string {
+func (DatabricksDialect) DataTypeForKind(kindDetails typing.KindDetails, _ bool) string {
 	switch kindDetails.Kind {
 	case typing.Float.Kind:
 		return "DOUBLE"

diff --git a/clients/databricks/dialect/typing_test.go b/clients/databricks/dialect/typing_test.go
@@ -0,0 +1,69 @@
+package dialect
+
+import (
+	"testing"
+
+	"github.com/artie-labs/transfer/lib/typing/decimal"
+
+	"github.com/artie-labs/transfer/lib/typing/ext"
+
+	"github.com/artie-labs/transfer/lib/typing"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDatabricksDialect_DataTypeForKind(t *testing.T) {
+	{
+		// Float
+		assert.Equal(t, "DOUBLE", DatabricksDialect{}.DataTypeForKind(typing.Float, false))
+	}
+	{
+		// Integer
+		assert.Equal(t, "BIGINT", DatabricksDialect{}.DataTypeForKind(typing.Integer, false))
+	}
+	{
+		// Variant
+		assert.Equal(t, "VARIANT", DatabricksDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.Struct.Kind}, false))
+	}
+	{
+		// Array
+		assert.Equal(t, "ARRAY<string>", DatabricksDialect{}.DataTypeForKind(typing.Array, false))
+	}
+	{
+		// String
+		assert.Equal(t, "STRING", DatabricksDialect{}.DataTypeForKind(typing.String, false))
+	}
+	{
+		// Boolean
+		assert.Equal(t, "BOOLEAN", DatabricksDialect{}.DataTypeForKind(typing.Boolean, false))
+	}
+	{
+		// Times
+		{
+			// Date
+			assert.Equal(t, "DATE", DatabricksDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.ETime.Kind, ExtendedTimeDetails: &ext.NestedKind{Type: ext.DateKindType}}, false))
+		}
+		{
+			// Timestamp
+			assert.Equal(t, "TIMESTAMP", DatabricksDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.ETime.Kind, ExtendedTimeDetails: &ext.NestedKind{Type: ext.TimestampTzKindType}}, false))
+		}
+		{
+			// Timestamp (w/o timezone)
+			assert.Equal(t, "TIMESTAMP", DatabricksDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.ETime.Kind, ExtendedTimeDetails: &ext.NestedKind{Type: ext.TimestampTzKindType}}, false))
+		}
+		{
+			// Time
+			assert.Equal(t, "STRING", DatabricksDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.ETime.Kind, ExtendedTimeDetails: &ext.NestedKind{Type: ext.TimeKindType}}, false))
+		}
+	}
+	{
+		// Decimals
+		{
+			// Below 38 precision
+			assert.Equal(t, "DECIMAL(10, 2)", DatabricksDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.EDecimal.Kind, ExtendedDecimalDetails: typing.ToPtr(decimal.NewDetails(10, 2))}, false))
+		}
+		{
+			// Above 38 precision
+			assert.Equal(t, "STRING", DatabricksDialect{}.DataTypeForKind(typing.KindDetails{Kind: typing.EDecimal.Kind, ExtendedDecimalDetails: typing.ToPtr(decimal.NewDetails(40, 2))}, false))
+		}
+	}
+}
diff --git a/lib/typing/decimal/base.go b/lib/typing/decimal/base.go
@@ -41,3 +41,11 @@ func (d Details) toKind(maxPrecision int32, exceededKind string) string {
 
 	return fmt.Sprintf("NUMERIC(%v, %v)", d.precision, d.scale)
 }
+
+func (d Details) toDecimalKind(maxPrecision int32, exceededKind string) string {
+	if d.precision > maxPrecision || d.precision == PrecisionNotSpecified {
+		return exceededKind
+	}
+
+	return fmt.Sprintf("DECIMAL(%d, %d)", d.precision, d.scale)
+}
diff --git a/lib/typing/decimal/details.go b/lib/typing/decimal/details.go
@@ -55,7 +55,7 @@ func (d Details) SnowflakeKind() string {
 // DatabricksKind - is used to determine whether a NUMERIC data type should be a STRING or NUMERIC(p, s).
 // Ref: https://docs.databricks.com/en/sql/language-manual/data-types/decimal-type.html
 func (d Details) DatabricksKind() string {
-	return d.toKind(MaxPrecisionBeforeString, "STRING")
+	return d.toDecimalKind(MaxPrecisionBeforeString, "STRING")
 }
 
 // MsSQLKind - Has the same limitation as Redshift