From 637c81f80ad5a4aeafce2a264d3f97736e000a1a Mon Sep 17 00:00:00 2001 From: Robin Tang Date: Wed, 20 Nov 2024 11:51:36 -0800 Subject: [PATCH] [DDL] Create table pass in columns (#1062) --- clients/bigquery/bigquery.go | 2 +- clients/databricks/store.go | 2 +- clients/mssql/staging.go | 2 +- clients/redshift/staging.go | 2 +- clients/shared/append.go | 2 +- clients/shared/merge.go | 2 +- clients/shared/table.go | 10 +++++++--- clients/snowflake/staging.go | 2 +- 8 files changed, 14 insertions(+), 10 deletions(-) diff --git a/clients/bigquery/bigquery.go b/clients/bigquery/bigquery.go index 5b3332506..9485b41f1 100644 --- 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 index 6713468d1..74d86b2c7 100644 --- 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 index 5bb2ebe2e..1a1272c07 100644 --- 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 index 99d465c77..6bc298ad8 100644 --- 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 index ec762064c..54f60d905 100644 --- 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 index 41ce87001..24aeb19d6 100644 --- 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 index aab6809a4..cc0a0dad1 100644 --- 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 index 4a09b6a12..c1ee740ab 100644 --- 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 } }