[dml] Split out Redshift delete query

artie-labs · May 3, 2024 · ca90900 · ca90900
1 parent c639579
commit ca90900
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 13 deletions.
diff --git a/lib/destination/dml/merge.go b/lib/destination/dml/merge.go
@@ -94,6 +94,19 @@ func (m *MergeArgument) buildRedshiftInsertQuery(columns []columns.Column, equal
 	)
 }
 
+func (m *MergeArgument) buildRedshiftDeleteQuery() string {
+	return fmt.Sprintf(`DELETE FROM %s WHERE (%s) IN (SELECT %s FROM %s as cc WHERE cc.%s = true);`,
+		// DELETE from table where (pk_1, pk_2)
+		m.TableID.FullyQualifiedName(), strings.Join(quoteColumns(m.PrimaryKeys, m.Dialect), ","),
+		// IN (cc.pk_1, cc.pk_2) FROM staging
+		array.StringsJoinAddPrefix(array.StringsJoinAddPrefixArgs{
+			Vals:      quoteColumns(m.PrimaryKeys, m.Dialect),
+			Separator: ",",
+			Prefix:    "cc.",
+		}), m.SubQuery, m.Dialect.QuoteIdentifier(constants.DeleteColumnMarker),
+	)
+}
+
 func (m *MergeArgument) GetParts() ([]string, error) {
 	if err := m.Valid(); err != nil {
 		return nil, err
@@ -161,18 +174,7 @@ func (m *MergeArgument) GetParts() ([]string, error) {
 	}
 
 	if *m.ContainsHardDeletes {
-		parts = append(parts,
-			// DELETE
-			fmt.Sprintf(`DELETE FROM %s WHERE (%s) IN (SELECT %s FROM %s as cc WHERE cc.%s = true);`,
-				// DELETE from table where (pk_1, pk_2)
-				m.TableID.FullyQualifiedName(), strings.Join(quoteColumns(m.PrimaryKeys, m.Dialect), ","),
-				// IN (cc.pk_1, cc.pk_2) FROM staging
-				array.StringsJoinAddPrefix(array.StringsJoinAddPrefixArgs{
-					Vals:      quoteColumns(m.PrimaryKeys, m.Dialect),
-					Separator: ",",
-					Prefix:    "cc.",
-				}), m.SubQuery, m.Dialect.QuoteIdentifier(constants.DeleteColumnMarker),
-			))
+		parts = append(parts, m.buildRedshiftDeleteQuery())
 	}
 
 	return parts, nil

diff --git a/lib/destination/dml/merge_test.go b/lib/destination/dml/merge_test.go
@@ -310,7 +310,7 @@ func TestMergeStatementEscapePrimaryKeys(t *testing.T) {
 	assert.Contains(t, mergeSQL, `cc."ID",cc."GROUP",cc."UPDATED_AT",cc."START"`, mergeSQL)
 }
 
-func TestBuildInsertQuery(t *testing.T) {
+func TestBuildRedshiftInsertQuery(t *testing.T) {
 	cols := []columns.Column{
 		columns.NewColumn("col1", typing.Invalid),
 		columns.NewColumn("col2", typing.Invalid),
@@ -327,3 +327,22 @@ func TestBuildInsertQuery(t *testing.T) {
 		mergeArg.buildRedshiftInsertQuery(cols, []string{"{EQUALITY_PART_1}", "{EQUALITY_PART_2}"}),
 	)
 }
+
+func TestBuildRedshiftDeleteQuery(t *testing.T) {
+	cols := []columns.Column{
+		columns.NewColumn("col1", typing.Invalid),
+		columns.NewColumn("col2", typing.Invalid),
+		columns.NewColumn("col3", typing.Invalid),
+	}
+
+	mergeArg := MergeArgument{
+		TableID:     MockTableIdentifier{"{TABLE_ID}"},
+		SubQuery:    "{SUB_QUERY}",
+		PrimaryKeys: []columns.Column{cols[0], cols[1]},
+		Dialect:     sql.SnowflakeDialect{},
+	}
+	assert.Equal(t,
+		`DELETE FROM {TABLE_ID} WHERE ("COL1","COL2") IN (SELECT cc."COL1",cc."COL2" FROM {SUB_QUERY} as cc WHERE cc."__ARTIE_DELETE" = true);`,
+		mergeArg.buildRedshiftDeleteQuery(),
+	)
+}