Improving column escape names (#1107)

Author: Tang8330

lib/typing/columns/columns.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"slices"
+	"strconv"
 	"strings"
 	"sync"
 
@@ -12,9 +13,16 @@ import (
 	"github.com/artie-labs/transfer/lib/typing"
 )
 
-// EscapeName - will lowercase columns and escape spaces.
 func EscapeName(name string) string {
+	// Lowercasing and escaping spaces.
 	_, name = stringutil.EscapeSpaces(strings.ToLower(name))
+
+	// Does the column name start with a number? If so, let's prefix `col_` to the column name.
+	// We're doing this most databases do not allow column names to start with a number.
+	if _, err := strconv.Atoi(string(name[0])); err == nil {
+		name = "col_" + name
+	}
+
 	return name
 }
 

lib/typing/columns/columns_test.go

@@ -29,6 +29,10 @@ func TestEscapeName(t *testing.T) {
 			name:         "col with spaces",
 			expectedName: "col__with__spaces",
 		},
+		{
+			name:         "1abc",
+			expectedName: "col_1abc",
+		},
 	}
 
 	for _, testCase := range testCases {

models/event/event.go

@@ -199,17 +199,10 @@ func (e *Event) Save(cfg config.Config, inMemDB *models.DatabaseData, tc kafkali
 	// Update col if necessary
 	sanitizedData := make(map[string]any)
 	for _col, val := range e.Data {
-		// TODO: Refactor this to call columns.EscapeName(...)
-		// columns need to all be normalized and lower cased.
-		newColName := strings.ToLower(_col)
-		// Columns here could contain spaces. Every destination treats spaces in a column differently.
-		// So far, Snowflake accepts them when escaped properly, however BigQuery does not accept it.
-		// Instead of making this more complicated for future destinations, we will escape the spaces by having double underscore `__`
-		// So, if customers want to retrieve spaces again, they can replace `__`.
-		var containsSpace bool
-		containsSpace, newColName = stringutil.EscapeSpaces(newColName)
-		if containsSpace {
-			// Write the message back if the column has changed.
+		newColName := columns.EscapeName(_col)
+		if newColName != _col {
+			// This means that the column name has changed.
+			// We need to update the column name in the sanitizedData map.
 			sanitizedData[newColName] = val
 		}
 

models/event/event_save_test.go

@@ -3,6 +3,7 @@ package event
 import (
 	"fmt"
 	"strconv"
+	"strings"
 
 	"github.com/artie-labs/transfer/lib/typing/columns"
 
@@ -173,21 +174,22 @@ func (e *EventsTestSuite) TestEventSaveOptionalSchema() {
 func (e *EventsTestSuite) TestEvent_SaveColumnsNoData() {
 	var cols columns.Columns
 	for i := 0; i < 50; i++ {
-		cols.AddColumn(columns.NewColumn(fmt.Sprint(i), typing.Invalid))
+		cols.AddColumn(columns.NewColumn(fmt.Sprintf("col_%d", i), typing.Invalid))
 	}
 
 	evt := Event{
 		Table:   "non_existent",
 		Columns: &cols,
 		Data: map[string]any{
-			"1":                                 "123",
+			"col_1":                             "123",
 			constants.DeleteColumnMarker:        true,
 			constants.OnlySetDeleteColumnMarker: true,
 		},
 		PrimaryKeyMap: map[string]any{
-			"1": "123",
+			"col_1": "123",
 		},
 	}
+
 	kafkaMsg := kafka.Message{}
 	_, _, err := evt.Save(e.cfg, e.db, topicConfig, artie.NewMessage(&kafkaMsg, nil, kafkaMsg.Topic))
 	assert.NoError(e.T(), err)
@@ -199,19 +201,21 @@ func (e *EventsTestSuite) TestEvent_SaveColumnsNoData() {
 			continue
 		}
 
+		columnNamePart := strings.Split(col.Name(), "_")[1]
+
 		if prevKey == "" {
-			prevKey = col.Name()
+			prevKey = columnNamePart
 			continue
 		}
 
-		currentKeyParsed, err := strconv.Atoi(col.Name())
+		currentKeyParsed, err := strconv.Atoi(columnNamePart)
 		assert.NoError(e.T(), err)
 
 		prevKeyParsed, err := strconv.Atoi(prevKey)
 		assert.NoError(e.T(), err)
 
 		// Testing ordering.
-		assert.True(e.T(), currentKeyParsed > prevKeyParsed, fmt.Sprintf("current key: %v, prevKey: %v", currentKeyParsed, prevKeyParsed))
+		assert.True(e.T(), currentKeyParsed > prevKeyParsed, fmt.Sprintf("current key: %q, prevKey: %q", currentKeyParsed, prevKeyParsed))
 	}
 
 	// Now let's add more keys.