[BigQuery] Use StorageWrite API + Dedupe (#726)

Co-authored-by: Nathan <148575555+nathan-artie@users.noreply.github.com>
artie-labs · Jun 17, 2024 · a320314 · a320314
1 parent 355beb5
commit a320314
Show file tree

Hide file tree

Showing 9 changed files with 34 additions and 11 deletions.
diff --git a/clients/bigquery/bigquery.go b/clients/bigquery/bigquery.go
@@ -34,6 +34,7 @@ const (
 	describeNameCol               = "column_name"
 	describeTypeCol               = "data_type"
 	describeCommentCol            = "description"
+	useStorageWriteAPI            = true
 )
 
 type Store struct {
@@ -71,7 +72,7 @@ func (s *Store) PrepareTemporaryTable(tableData *optimization.TableData, tableCo
 	}
 
 	// Load the data
-	if s.config.BigQuery.UseStorageWriteAPI {
+	if useStorageWriteAPI {
 		return s.putTableViaStorageWriteAPI(context.Background(), bqTempTableID, tableData)
 	} else {
 		return s.putTableViaLegacyAPI(context.Background(), bqTempTableID, tableData)

diff --git a/clients/bigquery/dialect/dialect.go b/clients/bigquery/dialect/dialect.go
@@ -123,7 +123,7 @@ func (BigQueryDialect) IsColumnAlreadyExistsErr(err error) bool {
 	return strings.Contains(err.Error(), "Column already exists")
 }
 
-func (BigQueryDialect) IsTableDoesNotExistErr(err error) bool {
+func (BigQueryDialect) IsTableDoesNotExistErr(_ error) bool {
 	return false
 }
 
@@ -154,6 +154,15 @@ func (bd BigQueryDialect) BuildIsNotToastValueExpression(tableAlias constants.Ta
 	return fmt.Sprintf("COALESCE(%s != '%s', true)", colName, constants.ToastUnavailableValuePlaceholder)
 }
 
+func (bd BigQueryDialect) BuildDedupeTableQuery(tableID sql.TableIdentifier, primaryKeys []string) string {
+	primaryKeysEscaped := sql.QuoteIdentifiers(primaryKeys, bd)
+	return fmt.Sprintf(`(SELECT * FROM %s QUALIFY ROW_NUMBER() OVER (PARTITION BY %s ORDER BY %s) = 1)`,
+		tableID.FullyQualifiedName(),
+		strings.Join(primaryKeysEscaped, ", "),
+		strings.Join(primaryKeysEscaped, ", "),
+	)
+}
+
 func (bd BigQueryDialect) BuildDedupeQueries(tableID, stagingTableID sql.TableIdentifier, primaryKeys []string, topicConfig kafkalib.TopicConfig) []string {
 	primaryKeysEscaped := sql.QuoteIdentifiers(primaryKeys, bd)
 

diff --git a/clients/bigquery/merge.go b/clients/bigquery/merge.go
@@ -51,6 +51,8 @@ func (s *Store) Merge(tableData *optimization.TableData) error {
 		AdditionalEqualityStrings: additionalEqualityStrings,
 		// BigQuery has DDL quotas.
 		RetryColBackfill: true,
+		// We are using BigQuery's streaming API which doesn't guarantee exactly once semantics
+		SubQueryDedupe: true,
 	})
 }
 

diff --git a/clients/mssql/dialect/dialect.go b/clients/mssql/dialect/dialect.go
@@ -169,6 +169,10 @@ func (md MSSQLDialect) BuildIsNotToastValueExpression(tableAlias constants.Table
 	return fmt.Sprintf("COALESCE(%s, '') != '%s'", colName, constants.ToastUnavailableValuePlaceholder)
 }
 
+func (MSSQLDialect) BuildDedupeTableQuery(tableID sql.TableIdentifier, primaryKeys []string) string {
+	panic("not implemented")
+}
+
 func (MSSQLDialect) BuildDedupeQueries(tableID, stagingTableID sql.TableIdentifier, primaryKeys []string, topicConfig kafkalib.TopicConfig) []string {
 	panic("not implemented") // We don't currently support deduping for MS SQL.
 }

diff --git a/clients/redshift/dialect/dialect.go b/clients/redshift/dialect/dialect.go
@@ -135,6 +135,10 @@ func (rd RedshiftDialect) BuildIsNotToastValueExpression(tableAlias constants.Ta
 	return fmt.Sprintf("COALESCE(%s != '%s', true)", colName, constants.ToastUnavailableValuePlaceholder)
 }
 
+func (rd RedshiftDialect) BuildDedupeTableQuery(tableID sql.TableIdentifier, _ []string) string {
+	return fmt.Sprintf(`( SELECT DISTINCT * FROM %s )`, tableID.FullyQualifiedName())
+}
+
 func (rd RedshiftDialect) BuildDedupeQueries(tableID, stagingTableID sql.TableIdentifier, primaryKeys []string, topicConfig kafkalib.TopicConfig) []string {
 	primaryKeysEscaped := sql.QuoteIdentifiers(primaryKeys, rd)
 

diff --git a/clients/shared/merge.go b/clients/shared/merge.go
@@ -111,10 +111,9 @@ func Merge(dwh destination.DataWarehouse, tableData *optimization.TableData, opt
 		}
 	}
 
-	temporaryTableName := temporaryTableID.FullyQualifiedName()
-	subQuery := temporaryTableName
+	subQuery := temporaryTableID.FullyQualifiedName()
 	if opts.SubQueryDedupe {
-		subQuery = fmt.Sprintf(`( SELECT DISTINCT * FROM %s )`, temporaryTableName)
+		subQuery = dwh.Dialect().BuildDedupeTableQuery(temporaryTableID, tableData.PrimaryKeys())
 	}
 
 	if subQuery == "" {

diff --git a/clients/snowflake/dialect/dialect.go b/clients/snowflake/dialect/dialect.go
@@ -153,6 +153,10 @@ func (sd SnowflakeDialect) BuildIsNotToastValueExpression(tableAlias constants.T
 	return fmt.Sprintf("COALESCE(%s != '%s', true)", colName, constants.ToastUnavailableValuePlaceholder)
 }
 
+func (SnowflakeDialect) BuildDedupeTableQuery(tableID sql.TableIdentifier, primaryKeys []string) string {
+	panic("not implemented")
+}
+
 func (sd SnowflakeDialect) BuildDedupeQueries(tableID, stagingTableID sql.TableIdentifier, primaryKeys []string, topicConfig kafkalib.TopicConfig) []string {
 	primaryKeysEscaped := sql.QuoteIdentifiers(primaryKeys, sd)
 

diff --git a/lib/config/bigquery.go b/lib/config/bigquery.go
@@ -5,12 +5,11 @@ import "fmt"
 type BigQuery struct {
 	// PathToCredentials is _optional_ if you have GOOGLE_APPLICATION_CREDENTIALS set as an env var
 	// Links to credentials: https://cloud.google.com/docs/authentication/application-default-credentials#GAC
-	PathToCredentials  string `yaml:"pathToCredentials"`
-	DefaultDataset     string `yaml:"defaultDataset"`
-	ProjectID          string `yaml:"projectID"`
-	Location           string `yaml:"location"`
-	BatchSize          int    `yaml:"batchSize"`
-	UseStorageWriteAPI bool   `yaml:"__useStorageWriteAPI"` // Not officially supported yet.
+	PathToCredentials string `yaml:"pathToCredentials"`
+	DefaultDataset    string `yaml:"defaultDataset"`
+	ProjectID         string `yaml:"projectID"`
+	Location          string `yaml:"location"`
+	BatchSize         int    `yaml:"batchSize"`
 }
 
 func (b *BigQuery) LoadDefaultValues() {

diff --git a/lib/sql/dialect.go b/lib/sql/dialect.go
@@ -24,6 +24,7 @@ type Dialect interface {
 	BuildCreateTableQuery(tableID TableIdentifier, temporary bool, colSQLParts []string) string
 	BuildAlterColumnQuery(tableID TableIdentifier, columnOp constants.ColumnOperation, colSQLPart string) string
 	BuildIsNotToastValueExpression(tableAlias constants.TableAlias, column columns.Column) string
+	BuildDedupeTableQuery(tableID TableIdentifier, primaryKeys []string) string
 	BuildDedupeQueries(tableID, stagingTableID TableIdentifier, primaryKeys []string, topicConfig kafkalib.TopicConfig) []string
 	BuildMergeQueries(
 		tableID TableIdentifier,