From c1f444ea54d6ab26fa5b91d17bf6b2e18d5ed5bf Mon Sep 17 00:00:00 2001
From: Nathan Villaescusa <nathan@artie.so>
Date: Fri, 3 May 2024 16:54:02 -0700
Subject: [PATCH] Update test

---
 lib/destination/dml/merge.go      |  8 +++---
 lib/destination/dml/merge_test.go | 42 +++++++++++++++++++++----------
 2 files changed, 33 insertions(+), 17 deletions(-)

diff --git a/lib/destination/dml/merge.go b/lib/destination/dml/merge.go
index a766498fb..94eb94daa 100644
--- a/lib/destination/dml/merge.go
+++ b/lib/destination/dml/merge.go
@@ -98,7 +98,7 @@ func (m *MergeArgument) buildRedshiftInsertQuery(columns []columns.Column) strin
 	)
 }
 
-func (m *MergeArgument) buildRedshiftUpdateQuery(columns []columns.Column, softDelete bool) string {
+func (m *MergeArgument) buildRedshiftUpdateQuery(columns []columns.Column) string {
 	stringBuilder := strings.Builder{}
 
 	stringBuilder.WriteString(fmt.Sprintf(`UPDATE %s as c SET %s FROM %s as cc WHERE %s`,
@@ -114,7 +114,7 @@ func (m *MergeArgument) buildRedshiftUpdateQuery(columns []columns.Column, softD
 		stringBuilder.WriteString(fmt.Sprintf(" AND cc.%s >= c.%s", m.IdempotentKey, m.IdempotentKey))
 	}
 
-	if !softDelete {
+	if !m.SoftDelete {
 		stringBuilder.WriteString(fmt.Sprintf(" AND COALESCE(cc.%s, false) = false", m.Dialect.QuoteIdentifier(constants.DeleteColumnMarker)))
 	}
 
@@ -158,7 +158,7 @@ func (m *MergeArgument) GetParts() ([]string, error) {
 	if m.SoftDelete {
 		return []string{
 			m.buildRedshiftInsertQuery(m.Columns),
-			m.buildRedshiftUpdateQuery(m.Columns, true),
+			m.buildRedshiftUpdateQuery(m.Columns),
 		}, nil
 	}
 
@@ -170,7 +170,7 @@ func (m *MergeArgument) GetParts() ([]string, error) {
 
 	parts := []string{
 		m.buildRedshiftInsertQuery(columns),
-		m.buildRedshiftUpdateQuery(columns, false),
+		m.buildRedshiftUpdateQuery(columns),
 	}
 
 	if *m.ContainsHardDeletes {
diff --git a/lib/destination/dml/merge_test.go b/lib/destination/dml/merge_test.go
index f1678e44c..76bf4749f 100644
--- a/lib/destination/dml/merge_test.go
+++ b/lib/destination/dml/merge_test.go
@@ -306,21 +306,37 @@ func TestMergeArgument_BuildRedshiftUpdateQuery(t *testing.T) {
 		columns.NewColumn("col3", typing.Invalid),
 	}
 
-	mergeArg := MergeArgument{
-		TableID:     MockTableIdentifier{"{TABLE_ID}"},
-		SubQuery:    "{SUB_QUERY}",
-		PrimaryKeys: []columns.Column{cols[0], cols[2]},
-		Dialect:     sql.SnowflakeDialect{},
+	{
+		// Soft deletes enabled:
+		mergeArg := MergeArgument{
+			TableID:     MockTableIdentifier{"{TABLE_ID}"},
+			SubQuery:    "{SUB_QUERY}",
+			PrimaryKeys: []columns.Column{cols[0], cols[2]},
+			Dialect:     sql.SnowflakeDialect{},
+			SoftDelete:  true,
+		}
+
+		assert.Equal(t,
+			`UPDATE {TABLE_ID} as c SET "COL1"=cc."COL1","COL2"=cc."COL2","COL3"=cc."COL3" FROM {SUB_QUERY} as cc WHERE c."COL1" = cc."COL1" and c."COL3" = cc."COL3";`,
+			mergeArg.buildRedshiftUpdateQuery(cols),
+		)
+	}
+	{
+		// Soft deletes disabled:
+		mergeArg := MergeArgument{
+			TableID:     MockTableIdentifier{"{TABLE_ID}"},
+			SubQuery:    "{SUB_QUERY}",
+			PrimaryKeys: []columns.Column{cols[0], cols[2]},
+			Dialect:     sql.SnowflakeDialect{},
+			SoftDelete:  false,
+		}
+
+		assert.Equal(t,
+			`UPDATE {TABLE_ID} as c SET "COL1"=cc."COL1","COL2"=cc."COL2","COL3"=cc."COL3" FROM {SUB_QUERY} as cc WHERE c."COL1" = cc."COL1" and c."COL3" = cc."COL3" AND COALESCE(cc."__ARTIE_DELETE", false) = false;`,
+			mergeArg.buildRedshiftUpdateQuery(cols),
+		)
 	}
 
-	assert.Equal(t,
-		`UPDATE {TABLE_ID} as c SET "COL1"=cc."COL1","COL2"=cc."COL2","COL3"=cc."COL3" FROM {SUB_QUERY} as cc WHERE c."COL1" = cc."COL1" and c."COL3" = cc."COL3";`,
-		mergeArg.buildRedshiftUpdateQuery(cols, true),
-	)
-	assert.Equal(t,
-		`UPDATE {TABLE_ID} as c SET "COL1"=cc."COL1","COL2"=cc."COL2","COL3"=cc."COL3" FROM {SUB_QUERY} as cc WHERE c."COL1" = cc."COL1" and c."COL3" = cc."COL3" AND COALESCE(cc."__ARTIE_DELETE", false) = false;`,
-		mergeArg.buildRedshiftUpdateQuery(cols, false),
-	)
 }
 
 func TestMergeArgument_BuildRedshiftDeleteQuery(t *testing.T) {