Merge pull request #85 from planetscale/dbussink/improve-type-mapping

Improve the type mapping
planetscale · Oct 17, 2023 · 40dce56 · 40dce56
2 parents d23b75a + a4878cb
commit 40dce56
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 41 deletions.
diff --git a/cmd/internal/planetscale_edge_database.go b/cmd/internal/planetscale_edge_database.go
@@ -128,35 +128,22 @@ func (p PlanetScaleEdgeDatabase) getStreamForTable(ctx context.Context, psc Plan
 func getJsonSchemaType(mysqlType string, treatTinyIntAsBoolean bool) PropertyType {
 	// Support custom airbyte types documented here :
 	// https://docs.airbyte.com/understanding-airbyte/supported-data-types/#the-types
-	if strings.HasPrefix(mysqlType, "int") {
-		return PropertyType{Type: "integer"}
-	}
-
-	if strings.HasPrefix(mysqlType, "decimal") || strings.HasPrefix(mysqlType, "double") {
-		return PropertyType{Type: "number"}
-	}
-
-	if strings.HasPrefix(mysqlType, "bigint") {
-		return PropertyType{Type: "string", AirbyteType: "big_integer"}
-	}
-
-	if strings.HasPrefix(mysqlType, "datetime") {
-		return PropertyType{Type: "string", CustomFormat: "date-time", AirbyteType: "timestamp_without_timezone"}
-	}
-
-	if strings.HasPrefix(mysqlType, "tinyint(1)") {
+	switch {
+	case strings.HasPrefix(mysqlType, "tinyint(1)"):
 		if treatTinyIntAsBoolean {
 			return PropertyType{Type: "boolean"}
 		}
-
-		return PropertyType{Type: "integer"}
-	}
-
-	switch mysqlType {
-	case "date":
-		return PropertyType{Type: "string", AirbyteType: "date"}
-	case "datetime":
-		return PropertyType{Type: "string", AirbyteType: "timestamp_without_timezone"}
+		return PropertyType{Type: "number", AirbyteType: "integer"}
+	case strings.HasPrefix(mysqlType, "int"), strings.HasPrefix(mysqlType, "smallint"), strings.HasPrefix(mysqlType, "mediumint"), strings.HasPrefix(mysqlType, "bigint"):
+		return PropertyType{Type: "number", AirbyteType: "integer"}
+	case strings.HasPrefix(mysqlType, "decimal"), strings.HasPrefix(mysqlType, "double"), strings.HasPrefix(mysqlType, "float"):
+		return PropertyType{Type: "number"}
+	case strings.HasPrefix(mysqlType, "datetime"), strings.HasPrefix(mysqlType, "timestamp"):
+		return PropertyType{Type: "string", CustomFormat: "date-time", AirbyteType: "timestamp_without_timezone"}
+	case strings.HasPrefix(mysqlType, "date"):
+		return PropertyType{Type: "string", CustomFormat: "date", AirbyteType: "date"}
+	case strings.HasPrefix(mysqlType, "time"):
+		return PropertyType{Type: "string", CustomFormat: "time", AirbyteType: "time_without_timezone"}
 	default:
 		return PropertyType{Type: "string"}
 	}

diff --git a/cmd/internal/planetscale_edge_database_test.go b/cmd/internal/planetscale_edge_database_test.go
@@ -4,11 +4,12 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"os"
+	"testing"
+
 	psdbconnect "github.com/planetscale/airbyte-source/proto/psdbconnect/v1alpha1"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
-	"os"
-	"testing"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/proto/query"
 )
@@ -209,9 +210,19 @@ func TestDiscover_CanPickRightAirbyteType(t *testing.T) {
 		TreatTinyIntAsBoolean bool
 	}{
 		{
-			MysqlType:      "int(32)",
-			JSONSchemaType: "integer",
-			AirbyteType:    "",
+			MysqlType:      "int(11)",
+			JSONSchemaType: "number",
+			AirbyteType:    "integer",
+		},
+		{
+			MysqlType:      "smallint(4)",
+			JSONSchemaType: "number",
+			AirbyteType:    "integer",
+		},
+		{
+			MysqlType:      "mediumint(8)",
+			JSONSchemaType: "number",
+			AirbyteType:    "integer",
 		},
 		{
 			MysqlType:             "tinyint(1)",
@@ -227,36 +238,51 @@ func TestDiscover_CanPickRightAirbyteType(t *testing.T) {
 		},
 		{
 			MysqlType:             "tinyint(1)",
-			JSONSchemaType:        "integer",
-			AirbyteType:           "",
+			JSONSchemaType:        "number",
+			AirbyteType:           "integer",
 			TreatTinyIntAsBoolean: false,
 		},
 		{
 			MysqlType:             "tinyint(1) unsigned",
-			JSONSchemaType:        "integer",
-			AirbyteType:           "",
+			JSONSchemaType:        "number",
+			AirbyteType:           "integer",
 			TreatTinyIntAsBoolean: false,
 		},
 		{
 			MysqlType:      "bigint(16)",
-			JSONSchemaType: "string",
-			AirbyteType:    "big_integer",
+			JSONSchemaType: "number",
+			AirbyteType:    "integer",
 		},
 		{
 			MysqlType:      "bigint unsigned",
-			JSONSchemaType: "string",
-			AirbyteType:    "big_integer",
+			JSONSchemaType: "number",
+			AirbyteType:    "integer",
 		},
 		{
 			MysqlType:      "bigint zerofill",
-			JSONSchemaType: "string",
-			AirbyteType:    "big_integer",
+			JSONSchemaType: "number",
+			AirbyteType:    "integer",
 		},
 		{
 			MysqlType:      "datetime",
 			JSONSchemaType: "string",
 			AirbyteType:    "timestamp_without_timezone",
 		},
+		{
+			MysqlType:      "datetime(6)",
+			JSONSchemaType: "string",
+			AirbyteType:    "timestamp_without_timezone",
+		},
+		{
+			MysqlType:      "time",
+			JSONSchemaType: "string",
+			AirbyteType:    "time_without_timezone",
+		},
+		{
+			MysqlType:      "time(6)",
+			JSONSchemaType: "string",
+			AirbyteType:    "time_without_timezone",
+		},
 		{
 			MysqlType:      "date",
 			JSONSchemaType: "string",
@@ -282,6 +308,11 @@ func TestDiscover_CanPickRightAirbyteType(t *testing.T) {
 			JSONSchemaType: "number",
 			AirbyteType:    "",
 		},
+		{
+			MysqlType:      "float(30)",
+			JSONSchemaType: "number",
+			AirbyteType:    "",
+		},
 	}
 
 	for _, typeTest := range tests {