[DDL] Refactor adding and dropping columns (#1052)

artie-labs · Nov 19, 2024 · 98abe01 · 98abe01
1 parent 8035d19
commit 98abe01
Show file tree

Hide file tree

Showing 17 changed files with 136 additions and 25 deletions.
diff --git a/clients/bigquery/dialect/ddl.go b/clients/bigquery/dialect/ddl.go
@@ -0,0 +1,14 @@
+package dialect
+
+import (
+	"github.com/artie-labs/transfer/lib/config/constants"
+	"github.com/artie-labs/transfer/lib/sql"
+)
+
+func (bd BigQueryDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string {
+	return bd.buildAlterColumnQuery(tableID, constants.Add, sqlPart)
+}
+
+func (bd BigQueryDialect) BuildDropColumnQuery(tableID sql.TableIdentifier, colName string) string {
+	return bd.buildAlterColumnQuery(tableID, constants.Delete, colName)
+}
diff --git a/clients/bigquery/dialect/dialect.go b/clients/bigquery/dialect/dialect.go
@@ -55,7 +55,7 @@ func (BigQueryDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, tempor
 	}
 }
 
-func (BigQueryDialect) BuildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
+func (BigQueryDialect) buildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
 	return fmt.Sprintf("ALTER TABLE %s %s COLUMN %s", tableID.FullyQualifiedName(), columnOp, colSQLPart)
 }
 

diff --git a/clients/bigquery/dialect/dialect_test.go b/clients/bigquery/dialect/dialect_test.go
@@ -66,13 +66,23 @@ func TestBigQueryDialect_BuildCreateTableQuery(t *testing.T) {
 	)
 }
 
-func TestBigQueryDialect_BuildAlterColumnQuery(t *testing.T) {
+func TestBigQueryDialect_BuildDropColumnQuery(t *testing.T) {
 	fakeTableID := &mocks.FakeTableIdentifier{}
 	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
 
 	assert.Equal(t,
 		"ALTER TABLE {TABLE} drop COLUMN {SQL_PART}",
-		BigQueryDialect{}.BuildAlterColumnQuery(fakeTableID, constants.Delete, "{SQL_PART}"),
+		BigQueryDialect{}.BuildDropColumnQuery(fakeTableID, "{SQL_PART}"),
+	)
+}
+
+func TestBigQueryDialect_BuildAddColumnQuery(t *testing.T) {
+	fakeTableID := &mocks.FakeTableIdentifier{}
+	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
+
+	assert.Equal(t,
+		"ALTER TABLE {TABLE} add COLUMN {SQL_PART}",
+		BigQueryDialect{}.BuildAddColumnQuery(fakeTableID, "{SQL_PART}"),
 	)
 }
 

diff --git a/clients/databricks/dialect/ddl.go b/clients/databricks/dialect/ddl.go
@@ -0,0 +1,14 @@
+package dialect
+
+import (
+	"github.com/artie-labs/transfer/lib/config/constants"
+	"github.com/artie-labs/transfer/lib/sql"
+)
+
+func (d DatabricksDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string {
+	return d.buildAlterColumnQuery(tableID, constants.Add, sqlPart)
+}
+
+func (d DatabricksDialect) BuildDropColumnQuery(tableID sql.TableIdentifier, colName string) string {
+	return d.buildAlterColumnQuery(tableID, constants.Delete, colName)
+}
diff --git a/clients/databricks/dialect/dialect.go b/clients/databricks/dialect/dialect.go
@@ -39,7 +39,7 @@ func (DatabricksDialect) BuildDescribeTableQuery(tableID sql.TableIdentifier) (s
 	return fmt.Sprintf("DESCRIBE TABLE %s", tableID.FullyQualifiedName()), nil, nil
 }
 
-func (DatabricksDialect) BuildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
+func (DatabricksDialect) buildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
 	return fmt.Sprintf("ALTER TABLE %s %s COLUMN %s", tableID.FullyQualifiedName(), columnOp, colSQLPart)
 }
 

diff --git a/clients/databricks/dialect/dialect_test.go b/clients/databricks/dialect/dialect_test.go
@@ -70,18 +70,16 @@ func TestDatabricksDialect_BuildCreateTableQuery(t *testing.T) {
 	}
 }
 
-func TestDatabricksDialect_BuildAlterColumnQuery(t *testing.T) {
+func TestDatabricksDialect_BuildAddColumnQuery(t *testing.T) {
 	fakeTableID := &mocks.FakeTableIdentifier{}
 	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
+	assert.Equal(t, "ALTER TABLE {TABLE} add COLUMN {SQL_PART} {DATA_TYPE}", DatabricksDialect{}.BuildAddColumnQuery(fakeTableID, "{SQL_PART} {DATA_TYPE}"))
+}
 
-	{
-		// DROP
-		assert.Equal(t, "ALTER TABLE {TABLE} drop COLUMN {SQL_PART}", DatabricksDialect{}.BuildAlterColumnQuery(fakeTableID, constants.Delete, "{SQL_PART}"))
-	}
-	{
-		// Add
-		assert.Equal(t, "ALTER TABLE {TABLE} add COLUMN {SQL_PART} {DATA_TYPE}", DatabricksDialect{}.BuildAlterColumnQuery(fakeTableID, constants.Add, "{SQL_PART} {DATA_TYPE}"))
-	}
+func TestDatabricksDialect_BuildDropColumnQuery(t *testing.T) {
+	fakeTableID := &mocks.FakeTableIdentifier{}
+	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
+	assert.Equal(t, "ALTER TABLE {TABLE} drop COLUMN {SQL_PART}", DatabricksDialect{}.BuildDropColumnQuery(fakeTableID, "{SQL_PART}"))
 }
 
 func TestDatabricksDialect_BuildDedupeQueries(t *testing.T) {

diff --git a/clients/mssql/dialect/ddl.go b/clients/mssql/dialect/ddl.go
@@ -0,0 +1,14 @@
+package dialect
+
+import (
+	"github.com/artie-labs/transfer/lib/config/constants"
+	"github.com/artie-labs/transfer/lib/sql"
+)
+
+func (md MSSQLDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string {
+	return md.buildAlterColumnQuery(tableID, constants.Add, sqlPart)
+}
+
+func (md MSSQLDialect) BuildDropColumnQuery(tableID sql.TableIdentifier, colName string) string {
+	return md.buildAlterColumnQuery(tableID, constants.Delete, colName)
+}
diff --git a/clients/mssql/dialect/dialect.go b/clients/mssql/dialect/dialect.go
@@ -72,7 +72,7 @@ func (MSSQLDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, _ bool, c
 	return fmt.Sprintf("CREATE TABLE %s (%s);", tableID.FullyQualifiedName(), strings.Join(colSQLParts, ","))
 }
 
-func (MSSQLDialect) BuildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
+func (MSSQLDialect) buildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
 	// Microsoft SQL Server doesn't support the COLUMN keyword
 	return fmt.Sprintf("ALTER TABLE %s %s %s", tableID.FullyQualifiedName(), columnOp, colSQLPart)
 }

diff --git a/clients/mssql/dialect/dialect_test.go b/clients/mssql/dialect/dialect_test.go
@@ -60,13 +60,23 @@ func TestMSSQLDialect_BuildCreateTableQuery(t *testing.T) {
 	)
 }
 
-func TestMSSQLDialect_BuildAlterColumnQuery(t *testing.T) {
+func TestMSSQLDialect_BuildAddColumnQuery(t *testing.T) {
+	fakeTableID := &mocks.FakeTableIdentifier{}
+	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
+
+	assert.Equal(t,
+		"ALTER TABLE {TABLE} add {SQL_PART}",
+		MSSQLDialect{}.BuildAddColumnQuery(fakeTableID, "{SQL_PART}"),
+	)
+}
+
+func TestMSSQLDialect_BuildDropColumnQuery(t *testing.T) {
 	fakeTableID := &mocks.FakeTableIdentifier{}
 	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
 
 	assert.Equal(t,
 		"ALTER TABLE {TABLE} drop {SQL_PART}",
-		MSSQLDialect{}.BuildAlterColumnQuery(fakeTableID, constants.Delete, "{SQL_PART}"),
+		MSSQLDialect{}.BuildDropColumnQuery(fakeTableID, "{SQL_PART}"),
 	)
 }
 

diff --git a/clients/redshift/dialect/ddl.go b/clients/redshift/dialect/ddl.go
@@ -0,0 +1,14 @@
+package dialect
+
+import (
+	"github.com/artie-labs/transfer/lib/config/constants"
+	"github.com/artie-labs/transfer/lib/sql"
+)
+
+func (rd RedshiftDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string {
+	return rd.buildAlterColumnQuery(tableID, constants.Add, sqlPart)
+}
+
+func (rd RedshiftDialect) BuildDropColumnQuery(tableID sql.TableIdentifier, colName string) string {
+	return rd.buildAlterColumnQuery(tableID, constants.Delete, colName)
+}
diff --git a/clients/redshift/dialect/dialect.go b/clients/redshift/dialect/dialect.go
@@ -35,7 +35,7 @@ func (RedshiftDialect) BuildCreateTableQuery(tableID sql.TableIdentifier, _ bool
 	return fmt.Sprintf("CREATE TABLE IF NOT EXISTS %s (%s);", tableID.FullyQualifiedName(), strings.Join(colSQLParts, ","))
 }
 
-func (RedshiftDialect) BuildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
+func (RedshiftDialect) buildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
 	return fmt.Sprintf("ALTER TABLE %s %s COLUMN %s", tableID.FullyQualifiedName(), columnOp, colSQLPart)
 }
 

diff --git a/clients/redshift/dialect/dialect_test.go b/clients/redshift/dialect/dialect_test.go
@@ -45,13 +45,23 @@ func TestRedshiftDialect_BuildCreateTableQuery(t *testing.T) {
 	)
 }
 
-func TestRedshiftDialect_BuildAlterColumnQuery(t *testing.T) {
+func TestRedshiftDialect_BuildAddColumnQuery(t *testing.T) {
+	fakeTableID := &mocks.FakeTableIdentifier{}
+	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
+
+	assert.Equal(t,
+		"ALTER TABLE {TABLE} add COLUMN {SQL_PART}",
+		RedshiftDialect{}.BuildAddColumnQuery(fakeTableID, "{SQL_PART}"),
+	)
+}
+
+func TestRedshiftDialect_BuildDropColumnQuery(t *testing.T) {
 	fakeTableID := &mocks.FakeTableIdentifier{}
 	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
 
 	assert.Equal(t,
 		"ALTER TABLE {TABLE} drop COLUMN {SQL_PART}",
-		RedshiftDialect{}.BuildAlterColumnQuery(fakeTableID, constants.Delete, "{SQL_PART}"),
+		RedshiftDialect{}.BuildDropColumnQuery(fakeTableID, "{SQL_PART}"),
 	)
 }
 

diff --git a/clients/snowflake/dialect/ddl.go b/clients/snowflake/dialect/ddl.go
@@ -0,0 +1,14 @@
+package dialect
+
+import (
+	"github.com/artie-labs/transfer/lib/config/constants"
+	"github.com/artie-labs/transfer/lib/sql"
+)
+
+func (sd SnowflakeDialect) BuildAddColumnQuery(tableID sql.TableIdentifier, sqlPart string) string {
+	return sd.buildAlterColumnQuery(tableID, constants.Add, sqlPart)
+}
+
+func (sd SnowflakeDialect) BuildDropColumnQuery(tableID sql.TableIdentifier, colName string) string {
+	return sd.buildAlterColumnQuery(tableID, constants.Delete, colName)
+}
diff --git a/clients/snowflake/dialect/dialect.go b/clients/snowflake/dialect/dialect.go
@@ -52,7 +52,7 @@ func (SnowflakeDialect) BuildDescribeTableQuery(tableID sql.TableIdentifier) (st
 	return fmt.Sprintf("DESC TABLE %s", tableID.FullyQualifiedName()), nil, nil
 }
 
-func (SnowflakeDialect) BuildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
+func (SnowflakeDialect) buildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
 	return fmt.Sprintf("ALTER TABLE %s %s COLUMN %s", tableID.FullyQualifiedName(), columnOp, colSQLPart)
 }
 

diff --git a/clients/snowflake/dialect/dialect_test.go b/clients/snowflake/dialect/dialect_test.go
@@ -57,13 +57,23 @@ func TestSnowflakeDialect_BuildCreateTableQuery(t *testing.T) {
 	)
 }
 
-func TestSnowflakeDialect_BuildAlterColumnQuery(t *testing.T) {
+func TestSnowflakeDialect_BuildAddColumnQuery(t *testing.T) {
+	fakeTableID := &mocks.FakeTableIdentifier{}
+	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
+
+	assert.Equal(t,
+		"ALTER TABLE {TABLE} add COLUMN {SQL_PART}",
+		SnowflakeDialect{}.BuildAddColumnQuery(fakeTableID, "{SQL_PART}"),
+	)
+}
+
+func TestSnowflakeDialect_BuildDropColumnQuery(t *testing.T) {
 	fakeTableID := &mocks.FakeTableIdentifier{}
 	fakeTableID.FullyQualifiedNameReturns("{TABLE}")
 
 	assert.Equal(t,
 		"ALTER TABLE {TABLE} drop COLUMN {SQL_PART}",
-		SnowflakeDialect{}.BuildAlterColumnQuery(fakeTableID, constants.Delete, "{SQL_PART}"),
+		SnowflakeDialect{}.BuildDropColumnQuery(fakeTableID, "{SQL_PART}"),
 	)
 }
 

diff --git a/lib/destination/ddl/ddl.go b/lib/destination/ddl/ddl.go
@@ -79,7 +79,7 @@ func BuildAlterTableAddColumns(dialect sql.Dialect, tableID sql.TableIdentifier,
 		}
 
 		sqlPart := fmt.Sprintf("%s %s", dialect.QuoteIdentifier(col.Name()), dialect.DataTypeForKind(col.KindDetails, col.PrimaryKey()))
-		parts = append(parts, dialect.BuildAlterColumnQuery(tableID, constants.Add, sqlPart))
+		parts = append(parts, dialect.BuildAddColumnQuery(tableID, sqlPart))
 	}
 
 	return parts, nil
@@ -90,5 +90,5 @@ func BuildAlterTableDropColumns(dialect sql.Dialect, tableID sql.TableIdentifier
 		return "", fmt.Errorf("received an invalid column %q", col.Name())
 	}
 
-	return dialect.BuildAlterColumnQuery(tableID, constants.Delete, dialect.QuoteIdentifier(col.Name())), nil
+	return dialect.BuildDropColumnQuery(tableID, dialect.QuoteIdentifier(col.Name())), nil
 }
diff --git a/lib/sql/dialect.go b/lib/sql/dialect.go
@@ -29,7 +29,6 @@ type Dialect interface {
 	KindForDataType(_type string, stringPrecision string) (typing.KindDetails, error)
 	IsColumnAlreadyExistsErr(err error) bool
 	IsTableDoesNotExistErr(err error) bool
-	BuildAlterColumnQuery(tableID TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string
 	BuildCreateTableQuery(tableID TableIdentifier, temporary bool, colSQLParts []string) string
 	BuildDedupeQueries(tableID, stagingTableID TableIdentifier, primaryKeys []string, includeArtieUpdatedAt bool) []string
 	BuildDedupeTableQuery(tableID TableIdentifier, primaryKeys []string) string
@@ -47,6 +46,10 @@ type Dialect interface {
 		containsHardDeletes bool,
 	) ([]string, error)
 
+	// DDL queries
+	BuildAddColumnQuery(tableID TableIdentifier, sqlPart string) string
+	BuildDropColumnQuery(tableID TableIdentifier, colName string) string
+
 	// Default values
 	GetDefaultValueStrategy() DefaultValueStrategy
 }
-Original file line number
+Diff line change
@@ Expand Up @@
     	}
     }
-    func (BigQueryDialect) BuildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
+    func (BigQueryDialect) buildAlterColumnQuery(tableID sql.TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string {
     	return fmt.Sprintf("ALTER TABLE %s %s COLUMN %s", tableID.FullyQualifiedName(), columnOp, colSQLPart)
     }
@@ Expand Down @@