materialize-snowflake: transient tables + split insert and merge

Author: mdibaiee

materialize-snowflake/snowflake.go

@@ -19,6 +19,7 @@ import (
 	sql "github.com/estuary/connectors/materialize-sql"
 	pf "github.com/estuary/flow/go/protocols/flow"
 	pm "github.com/estuary/flow/go/protocols/materialize"
+	"github.com/google/uuid"
 	"github.com/jmoiron/sqlx"
 	log "github.com/sirupsen/logrus"
 	sf "github.com/snowflakedb/gosnowflake"
@@ -326,10 +327,13 @@ type binding struct {
 	}
 	// Variables accessed by Prepare, Store, and Commit.
 	store struct {
-		stage       *stagedFile
+		insertStage *stagedFile // For rows where !it.Exists (new rows)
+		mergeStage  *stagedFile // For rows where it.Exists (existing rows)
+		hasInserts  bool        // Track if we have insert operations
+		hasMerges   bool        // Track if we have merge operations
 		mergeInto   string
 		copyInto    string
-		mustMerge   bool
+		mustMerge   bool // Deprecated: use hasMerges instead
 		mergeBounds *sql.MergeBoundsBuilder
 	}
 }
@@ -365,7 +369,8 @@ func (d *transactor) addBinding(ctx context.Context, target sql.Table, streaming
 	}
 
 	b.load.stage = newStagedFile(os.TempDir())
-	b.store.stage = newStagedFile(os.TempDir())
+	b.store.insertStage = newStagedFile(os.TempDir())
+	b.store.mergeStage = newStagedFile(os.TempDir())
 
 	if b.target.DeltaUpdates && d.cfg.Credentials.AuthType == snowflake_auth.JWT {
 		var keyBegin = fmt.Sprintf("%08x", d._range.KeyBegin)
@@ -557,8 +562,14 @@ func (d *transactor) pipeExists(ctx context.Context, pipeName string) (bool, err
 
 type checkpointItem struct {
 	Table       string
-	Query       string
-	StagedDir   string
+	Query       string // Deprecated: use InsertQuery/MergeQuery for new code
+	InsertQuery string // Query for insert operations (from insertStage)
+	MergeQuery  string // Query for merge operations (from mergeStage)
+	StagedDir   string // Deprecated: use InsertStagedDir/MergeStagedDir
+	InsertStagedDir string // Staged directory for inserts
+	MergeStagedDir  string // Staged directory for merges
+	InsertStagingTable string // Transient table name for inserts (for cleanup)
+	MergeStagingTable  string // Transient table name for merges (for cleanup)
 	StreamBlobs []*blobMetadata
 	PipeName    string
 	PipeFiles   []fileRecord
@@ -579,12 +590,21 @@ func (d *transactor) Store(it *m.StoreIterator) (m.StartCommitFunc, error) {
 			continue
 		}
 
+		// Route to appropriate stage based on whether key was loaded
+		var targetStage *stagedFile
 		if it.Exists {
-			b.store.mustMerge = true
+			// Key was loaded, so it exists - this is an UPDATE or DELETE
+			targetStage = b.store.mergeStage
+			b.store.hasMerges = true
+			b.store.mustMerge = true // Keep for backward compatibility
+		} else {
+			// Key was not loaded, so it doesn't exist - this is an INSERT
+			targetStage = b.store.insertStage
+			b.store.hasInserts = true
 		}
 
 		if !b.streaming {
-			if err := b.store.stage.start(ctx, d.db); err != nil {
+			if err := targetStage.start(ctx, d.db); err != nil {
 				return nil, err
 			}
 		}
@@ -594,10 +614,13 @@ func (d *transactor) Store(it *m.StoreIterator) (m.StartCommitFunc, error) {
 			if err := d.streamManager.writeRow(ctx, it.Binding, converted); err != nil {
 				return nil, fmt.Errorf("encoding Store to stream for resource %s: %w", b.target.Path, err)
 			}
-		} else if err = b.store.stage.writeRow(append(converted, flowDelete)); err != nil {
+		} else if err = targetStage.writeRow(append(converted, flowDelete)); err != nil {
 			return nil, fmt.Errorf("writing Store to scratch file: %w", err)
 		} else {
-			b.store.mergeBounds.NextKey(converted[:len(b.target.Keys)])
+			// Only track merge bounds for rows that will be merged
+			if it.Exists {
+				b.store.mergeBounds.NextKey(converted[:len(b.target.Keys)])
+			}
 		}
 	}
 	if it.Err() != nil {
@@ -638,37 +661,84 @@ func (d *transactor) buildDriverCheckpoint(ctx context.Context, runtimeCheckpoin
 			continue
 		}
 
-		if !b.store.stage.started {
+		// Skip if neither stage has data
+		if !b.store.insertStage.started && !b.store.mergeStage.started {
 			continue
 		}
 
-		dir, err := b.store.stage.flush()
-		if err != nil {
-			return nil, err
+		var cpItem = &checkpointItem{
+			Table:   b.target.Identifier,
+			Version: d.version,
 		}
 
-		if b.store.mustMerge {
-			mergeIntoQuery, err := renderBoundedQueryTemplate(d.templates.mergeInto, b.target, dir, b.store.mergeBounds.Build())
+		// Handle INSERT operations (new rows)
+		if b.store.hasInserts && b.store.insertStage.started {
+			insertDir, err := b.store.insertStage.flush()
 			if err != nil {
-				return nil, fmt.Errorf("mergeInto template: %w", err)
+				return nil, fmt.Errorf("flushing insert stage: %w", err)
 			}
-			d.cp[b.target.StateKey] = &checkpointItem{
-				Table:     b.target.Identifier,
-				Query:     mergeIntoQuery,
-				StagedDir: dir,
-				Version:   d.version,
+
+			if b.pipeName != "" {
+				// Snowpipe for inserts
+				cpItem.InsertStagedDir = insertDir
+				cpItem.PipeFiles = b.store.insertStage.uploaded
+				cpItem.PipeName = b.pipeName
+			} else {
+				// Direct COPY INTO for inserts (no transient table needed)
+				insertQuery, err := renderTableAndFileTemplate(b.target, insertDir, d.templates.copyInto)
+				if err != nil {
+					return nil, fmt.Errorf("copyInto template for inserts: %w", err)
+				}
+				cpItem.InsertQuery = insertQuery
+				cpItem.InsertStagedDir = insertDir
 			}
-			// Reset for next round.
-			b.store.mustMerge = false
-		} else if b.pipeName != "" {
+		}
+
+		// Handle MERGE operations (existing rows - updates/deletes)
+		if b.store.hasMerges && b.store.mergeStage.started {
+			mergeDir, err := b.store.mergeStage.flush()
+			if err != nil {
+				return nil, fmt.Errorf("flushing merge stage: %w", err)
+			}
+
+			// Generate unique transient table name for merges
+			// Replace hyphens with underscores to avoid identifier quoting issues
+			mergeStagingTable := fmt.Sprintf("%s.FLOW_STAGING_MERGE_%d_%s",
+				d.cfg.Schema, idx, strings.ReplaceAll(uuid.New().String(), "-", "_"))
+
+			// Build compound query: CREATE + COPY + MERGE + DROP
+			createSQL, err := renderCreateTransientTable(b.target, mergeStagingTable, d.templates.createTransientTable)
+			if err != nil {
+				return nil, fmt.Errorf("createTransientTable template for merges: %w", err)
+			}
+
+			copySQL, err := renderCopyIntoTransient(b.target, mergeStagingTable, mergeDir, d.templates.copyIntoTransient)
+			if err != nil {
+				return nil, fmt.Errorf("copyIntoTransient template for merges: %w", err)
+			}
+
+			mergeSQL, err := renderMergeFromTransient(b.target, mergeStagingTable, b.store.mergeBounds.Build(), d.templates.mergeFromTransient)
+			if err != nil {
+				return nil, fmt.Errorf("mergeFromTransient template: %w", err)
+			}
+
+			dropSQL := fmt.Sprintf("DROP TABLE IF EXISTS %s;", mergeStagingTable)
+
+			cpItem.MergeQuery = fmt.Sprintf("%s\n%s\n%s\n%s",
+				createSQL, copySQL, mergeSQL, dropSQL)
+			cpItem.MergeStagedDir = mergeDir
+			cpItem.MergeStagingTable = mergeStagingTable
+		}
+
+		// Handle Snowpipe creation if needed
+		if b.pipeName != "" && (b.store.hasInserts || b.store.hasMerges) {
 			// Check to see if a pipe for this version already exists
 			exists, err := d.pipeExists(ctx, b.pipeName)
 			if err != nil {
 				return nil, err
 			}
 
-			// Only create the pipe if it doesn't exist. Since the pipe name is versioned by the spec
-			// it means if the spec has been updated, we will end up creating a new pipe
+			// Only create the pipe if it doesn't exist
 			if !exists {
 				log.WithField("name", b.pipeName).Info("store: creating pipe")
 				if createPipe, err := renderTablePipeTemplate(b.target, b.pipeName, d.templates.createPipe); err != nil {
@@ -678,8 +748,7 @@ func (d *transactor) buildDriverCheckpoint(ctx context.Context, runtimeCheckpoin
 				}
 			}
 
-			// Our understanding is that CREATE PIPE is _eventually consistent_, and so we
-			// wait until we can make sure the pipe exists before continuing
+			// Wait for pipe to be eventually consistent
 			for !exists {
 				exists, err = d.pipeExists(ctx, b.pipeName)
 				if err != nil {
@@ -689,25 +758,15 @@ func (d *transactor) buildDriverCheckpoint(ctx context.Context, runtimeCheckpoin
 					time.Sleep(5 * time.Second)
 				}
 			}
+		}
 
-			d.cp[b.target.StateKey] = &checkpointItem{
-				Table:     b.target.Identifier,
-				StagedDir: dir,
-				PipeFiles: b.store.stage.uploaded,
-				PipeName:  b.pipeName,
-				Version:   d.version,
-			}
-
-		} else {
-			if copyIntoQuery, err := renderTableAndFileTemplate(b.target, dir, d.templates.copyInto); err != nil {
-				return nil, fmt.Errorf("copyInto template: %w", err)
-			} else {
-				d.cp[b.target.StateKey] = &checkpointItem{
-					Table:     b.target.Identifier,
-					Query:     copyIntoQuery,
-					StagedDir: dir,
-				}
-			}
+		// Store checkpoint item if we have any operations
+		if cpItem.InsertQuery != "" || cpItem.MergeQuery != "" || cpItem.PipeName != "" {
+			d.cp[b.target.StateKey] = cpItem
+			// Reset flags for next transaction
+			b.store.hasInserts = false
+			b.store.hasMerges = false
+			b.store.mustMerge = false
 		}
 	}
 
@@ -808,25 +867,112 @@ func (d *transactor) Acknowledge(ctx context.Context) (*pf.ConnectorState, error
 			continue
 		}
 
+		// Execute INSERT query if present
+		if len(item.InsertQuery) > 0 {
+			item := item
+			group.Go(func() error {
+				d.be.StartedResourceCommit(path)
+
+				queryStart := time.Now()
+				result, err := d.db.ExecContext(ctx, item.InsertQuery)
+				queryDuration := time.Since(queryStart)
+
+				if err != nil {
+					// Best-effort cleanup of transient table on failure
+					d.cleanupTransientTable(ctx, item.InsertStagingTable)
+					return fmt.Errorf("insert query failed: %w", err)
+				}
+
+				// Log performance metrics for performance test tables
+				if strings.HasPrefix(item.Table, "perf_") {
+					rowsAffected, _ := result.RowsAffected()
+					log.WithFields(log.Fields{
+						"perf_query_duration_ms": queryDuration.Milliseconds(),
+						"perf_rows_affected":     rowsAffected,
+						"perf_table":             item.Table,
+						"perf_query_type":        "INSERT",
+					}).Info("SNOWFLAKE_PERF_METRIC")
+				}
+
+				d.be.FinishedResourceCommit(path)
+				if err := d.deleteFiles(groupCtx, []string{item.InsertStagedDir}); err != nil {
+					return fmt.Errorf("cleaning up insert files: %w", err)
+				}
+
+				return nil
+			})
+		}
+
+		// Execute MERGE query if present
+		if len(item.MergeQuery) > 0 {
+			item := item
+			group.Go(func() error {
+				d.be.StartedResourceCommit(path)
+
+				queryStart := time.Now()
+				result, err := d.db.ExecContext(ctx, item.MergeQuery)
+				queryDuration := time.Since(queryStart)
+
+				if err != nil {
+					// Best-effort cleanup of transient table on failure
+					d.cleanupTransientTable(ctx, item.MergeStagingTable)
+					return fmt.Errorf("merge query failed: %w", err)
+				}
+
+				// Log performance metrics for performance test tables
+				if strings.HasPrefix(item.Table, "perf_") {
+					rowsAffected, _ := result.RowsAffected()
+					log.WithFields(log.Fields{
+						"perf_query_duration_ms": queryDuration.Milliseconds(),
+						"perf_rows_affected":     rowsAffected,
+						"perf_table":             item.Table,
+						"perf_query_type":        "MERGE",
+					}).Info("SNOWFLAKE_PERF_METRIC")
+				}
+
+				d.be.FinishedResourceCommit(path)
+				if err := d.deleteFiles(groupCtx, []string{item.MergeStagedDir}); err != nil {
+					return fmt.Errorf("cleaning up merge files: %w", err)
+				}
+
+				return nil
+			})
+		}
+
+		// Legacy: Handle old Query field for backward compatibility
 		if len(item.Query) > 0 {
 			item := item
 			group.Go(func() error {
 				d.be.StartedResourceCommit(path)
-				// NB: Not using groupTx here since the Go Snowflake driver
-				// retains contexts internally, and groupCtx is cancelled after
-				// group.Wait() returns.
-				if _, err := d.db.ExecContext(ctx, item.Query); err != nil {
+
+				queryStart := time.Now()
+				result, err := d.db.ExecContext(ctx, item.Query)
+				queryDuration := time.Since(queryStart)
+
+				if err != nil {
 					return fmt.Errorf("query %q failed: %w", item.Query, err)
 				}
 
+				// Log performance metrics for performance test tables
+				if strings.HasPrefix(item.Table, "perf_") {
+					rowsAffected, _ := result.RowsAffected()
+					log.WithFields(log.Fields{
+						"perf_query_duration_ms": queryDuration.Milliseconds(),
+						"perf_rows_affected":     rowsAffected,
+						"perf_table":             item.Table,
+					}).Info("SNOWFLAKE_PERF_METRIC")
+				}
+
 				d.be.FinishedResourceCommit(path)
 				if err := d.deleteFiles(groupCtx, []string{item.StagedDir}); err != nil {
 					return fmt.Errorf("cleaning up files: %w", err)
 				}
 
 				return nil
 			})
-		} else if len(item.StreamBlobs) > 0 {
+		}
+
+		if len(item.StreamBlobs) > 0 {
 			group.Go(func() error {
 				d.be.StartedResourceCommit(path)
 				if err := d.streamManager.write(groupCtx, item.StreamBlobs); err != nil {
@@ -1048,6 +1194,19 @@ func (d *transactor) Acknowledge(ctx context.Context) (*pf.ConnectorState, error
 	return &pf.ConnectorState{UpdatedJson: json.RawMessage(checkpointJSON), MergePatch: true}, nil
 }
 
+// Best-effort cleanup of transient staging tables on error
+func (d *transactor) cleanupTransientTable(ctx context.Context, tableName string) {
+	if tableName == "" {
+		return
+	}
+	dropSQL := fmt.Sprintf("DROP TABLE IF EXISTS %s", tableName)
+	if _, err := d.db.ExecContext(ctx, dropSQL); err != nil {
+		log.WithError(err).WithField("table", tableName).Warn("failed to cleanup transient staging table")
+	} else {
+		log.WithField("table", tableName).Debug("cleaned up transient staging table")
+	}
+}
+
 func (d *transactor) pathForStateKey(stateKey string) []string {
 	for _, b := range d.bindings {
 		if b.target.StateKey == stateKey {