Merge branch 'master' into disable-time-travel-for-temp-tables

artie-labs · Nov 20, 2024 · b423ea1 · b423ea1
2 parents 1e3b901 + 637c81f
commit b423ea1
Show file tree

Hide file tree

Showing 8 changed files with 14 additions and 10 deletions.
diff --git a/clients/bigquery/bigquery.go b/clients/bigquery/bigquery.go
@@ -82,7 +82,7 @@ func (s *Store) Append(ctx context.Context, tableData *optimization.TableData, u
 
 func (s *Store) PrepareTemporaryTable(ctx context.Context, tableData *optimization.TableData, dwh *types.DwhTableConfig, tempTableID sql.TableIdentifier, _ sql.TableIdentifier, opts types.AdditionalSettings, createTempTable bool) error {
 	if createTempTable {
-		if err := shared.CreateTable(ctx, s, tableData, dwh, opts.ColumnSettings, tempTableID, true); err != nil {
+		if err := shared.CreateTempTable(ctx, s, tableData, dwh, opts.ColumnSettings, tempTableID); err != nil {
 			return err
 		}
 	}

diff --git a/clients/databricks/store.go b/clients/databricks/store.go
@@ -82,7 +82,7 @@ func (s Store) GetTableConfig(tableData *optimization.TableData) (*types.DwhTabl
 
 func (s Store) PrepareTemporaryTable(ctx context.Context, tableData *optimization.TableData, dwh *types.DwhTableConfig, tempTableID sql.TableIdentifier, _ sql.TableIdentifier, opts types.AdditionalSettings, createTempTable bool) error {
 	if createTempTable {
-		if err := shared.CreateTable(ctx, s, tableData, dwh, opts.ColumnSettings, tempTableID, true); err != nil {
+		if err := shared.CreateTempTable(ctx, s, tableData, dwh, opts.ColumnSettings, tempTableID); err != nil {
 			return err
 		}
 	}

diff --git a/clients/mssql/staging.go b/clients/mssql/staging.go
@@ -15,7 +15,7 @@ import (
 
 func (s *Store) PrepareTemporaryTable(ctx context.Context, tableData *optimization.TableData, dwh *types.DwhTableConfig, tempTableID sql.TableIdentifier, _ sql.TableIdentifier, opts types.AdditionalSettings, createTempTable bool) error {
 	if createTempTable {
-		if err := shared.CreateTable(ctx, s, tableData, dwh, opts.ColumnSettings, tempTableID, true); err != nil {
+		if err := shared.CreateTempTable(ctx, s, tableData, dwh, opts.ColumnSettings, tempTableID); err != nil {
 			return err
 		}
 	}

diff --git a/clients/redshift/staging.go b/clients/redshift/staging.go
@@ -40,7 +40,7 @@ func (s *Store) PrepareTemporaryTable(ctx context.Context, tableData *optimizati
 	}
 
 	if createTempTable {
-		if err = shared.CreateTable(ctx, s, tableData, tableConfig, opts.ColumnSettings, tempTableID, true); err != nil {
+		if err = shared.CreateTempTable(ctx, s, tableData, tableConfig, opts.ColumnSettings, tempTableID); err != nil {
 			return err
 		}
 	}

diff --git a/clients/shared/append.go b/clients/shared/append.go
@@ -32,7 +32,7 @@ func Append(ctx context.Context, dwh destination.DataWarehouse, tableData *optim
 
 	tableID := dwh.IdentifierFor(tableData.TopicConfig(), tableData.Name())
 	if tableConfig.CreateTable() {
-		if err = CreateTable(ctx, dwh, tableData, tableConfig, opts.ColumnSettings, tableID, false); err != nil {
+		if err = CreateTable(ctx, dwh, tableData, tableConfig, opts.ColumnSettings, tableID, false, targetKeysMissing); err != nil {
 			return fmt.Errorf("failed to create table: %w", err)
 		}
 	} else {

diff --git a/clients/shared/merge.go b/clients/shared/merge.go
@@ -38,7 +38,7 @@ func Merge(ctx context.Context, dwh destination.DataWarehouse, tableData *optimi
 
 	tableID := dwh.IdentifierFor(tableData.TopicConfig(), tableData.Name())
 	if tableConfig.CreateTable() {
-		if err = CreateTable(ctx, dwh, tableData, tableConfig, opts.ColumnSettings, tableID, false); err != nil {
+		if err = CreateTable(ctx, dwh, tableData, tableConfig, opts.ColumnSettings, tableID, false, targetKeysMissing); err != nil {
 			return fmt.Errorf("failed to create table: %w", err)
 		}
 	} else {

diff --git a/clients/shared/table.go b/clients/shared/table.go
@@ -16,8 +16,12 @@ import (
 	"github.com/artie-labs/transfer/lib/typing/columns"
 )
 
-func CreateTable(ctx context.Context, dwh destination.DataWarehouse, tableData *optimization.TableData, tc *types.DwhTableConfig, settings config.SharedDestinationColumnSettings, tableID sql.TableIdentifier, tempTable bool) error {
-	query, err := ddl.BuildCreateTableSQL(settings, dwh.Dialect(), tableID, tempTable, tableData.Mode(), tableData.ReadOnlyInMemoryCols().GetColumns())
+func CreateTempTable(ctx context.Context, dwh destination.DataWarehouse, tableData *optimization.TableData, tc *types.DwhTableConfig, settings config.SharedDestinationColumnSettings, tableID sql.TableIdentifier) error {
+	return CreateTable(ctx, dwh, tableData, tc, settings, tableID, true, tableData.ReadOnlyInMemoryCols().GetColumns())
+}
+
+func CreateTable(ctx context.Context, dwh destination.DataWarehouse, tableData *optimization.TableData, tc *types.DwhTableConfig, settings config.SharedDestinationColumnSettings, tableID sql.TableIdentifier, tempTable bool, cols []columns.Column) error {
+	query, err := ddl.BuildCreateTableSQL(settings, dwh.Dialect(), tableID, tempTable, tableData.Mode(), cols)
 	if err != nil {
 		return fmt.Errorf("failed to build create table sql: %w", err)
 	}
@@ -28,7 +32,7 @@ func CreateTable(ctx context.Context, dwh destination.DataWarehouse, tableData *
 	}
 
 	// Update cache with the new columns that we've added.
-	tc.MutateInMemoryColumns(constants.Add, tableData.ReadOnlyInMemoryCols().GetColumns()...)
+	tc.MutateInMemoryColumns(constants.Add, cols...)
 	return nil
 }
 

diff --git a/clients/snowflake/staging.go b/clients/snowflake/staging.go
@@ -52,7 +52,7 @@ func castColValStaging(colVal any, colKind typing.KindDetails) (string, error) {
 
 func (s *Store) PrepareTemporaryTable(ctx context.Context, tableData *optimization.TableData, dwh *types.DwhTableConfig, tempTableID sql.TableIdentifier, _ sql.TableIdentifier, additionalSettings types.AdditionalSettings, createTempTable bool) error {
 	if createTempTable {
-		if err := shared.CreateTable(ctx, s, tableData, dwh, additionalSettings.ColumnSettings, tempTableID, true); err != nil {
+		if err := shared.CreateTempTable(ctx, s, tableData, dwh, additionalSettings.ColumnSettings, tempTableID); err != nil {
 			return err
 		}
 	}