Cherry pick #450 repl: track updates to PRIMARY KEY correctly

Author: morgo

pkg/repl/client.go

@@ -387,31 +387,77 @@ func (c *Client) processRowsEvent(ev *replication.BinlogEvent, e *replication.Ro
 		return nil // ignore event, it could be to a _new table.
 	}
 	eventType := parseEventType(ev.Header.EventType)
-	var i = 0
-	for _, row := range e.Rows {
-		if eventType == eventTypeUpdate {
-			// For update events there are always before and after images (i.e. e.Rows is always in pairs.)
-			// We only need to capture one of the events, and since in MINIMAL RBR row
-			// image the PK is only included in the before, we chose that one.
-			i++
-			if i%2 == 0 {
-				continue
+
+	if eventType == eventTypeUpdate {
+		// For update events there are always before and after images (i.e. e.Rows is always in pairs.)
+		// With MINIMAL row image, the PK is only included in the before image for non-PK updates.
+		// For PK updates, both before and after images will contain the PK columns since they changed.
+		for i := 0; i < len(e.Rows); i += 2 {
+			beforeRow := e.Rows[i]
+			afterRow := e.Rows[i+1]
+
+			// Always process the before image (guaranteed to have PK in minimal mode)
+			beforeKey, err := sub.table.PrimaryKeyValues(beforeRow)
+			if err != nil {
+				return err
+			}
+			if len(beforeKey) == 0 {
+				return fmt.Errorf("no primary key found for before row: %#v", beforeRow)
+			}
+
+			// With MINIMAL row image, we need to reconstruct the after key
+			// by combining the before key with any changed PK columns from the after image
+			afterKey := make([]any, len(beforeKey))
+			copy(afterKey, beforeKey) // Start with the before key
+
+			// Check if any PK columns were updated by examining the after row
+			isPKUpdate := false
+			afterRowSlice := afterRow
+
+			for pkIdx, pkCol := range sub.table.KeyColumns {
+				// Find the position of this PK column in the table columns
+				for colIdx, col := range sub.table.Columns {
+					if col == pkCol {
+						// If this column exists in the after image and is not nil, use it
+						if colIdx < len(afterRowSlice) && afterRowSlice[colIdx] != nil {
+							if fmt.Sprintf("%v", beforeKey[pkIdx]) != fmt.Sprintf("%v", afterRowSlice[colIdx]) {
+								afterKey[pkIdx] = afterRowSlice[colIdx]
+								isPKUpdate = true
+							}
+						}
+						break
+					}
+				}
+			}
+
+			if isPKUpdate {
+				// This is a primary key update - track both delete and insert
+				sub.keyHasChanged(beforeKey, true)  // delete old key
+				sub.keyHasChanged(afterKey, false) // insert new key
+			} else {
+				// Same PK, just a regular update
+				sub.keyHasChanged(beforeKey, false)
 			}
 		}
-		key, err := sub.table.PrimaryKeyValues(row)
-		if err != nil {
-			return err
-		}
-		if len(key) == 0 {
-			return fmt.Errorf("no primary key found for row: %#v", row)
-		}
-		switch eventType {
-		case eventTypeInsert, eventTypeUpdate:
-			sub.keyHasChanged(key, false)
-		case eventTypeDelete:
-			sub.keyHasChanged(key, true)
-		default:
-			c.logger.Errorf("unknown event type: %v", ev.Header.EventType)
+	} else {
+		// For INSERT and DELETE events, process each row normally
+		for _, row := range e.Rows {
+			key, err := sub.table.PrimaryKeyValues(row)
+			if err != nil {
+				return err
+			}
+			if len(key) == 0 {
+				// In theory this is unreachable since we mandate a PK on tables
+				return fmt.Errorf("no primary key found for row: %#v", row)
+			}
+			switch eventType {
+			case eventTypeInsert:
+				sub.keyHasChanged(key, false)
+			case eventTypeDelete:
+				sub.keyHasChanged(key, true)
+			default:
+				c.logger.Errorf("unknown event type: %v", ev.Header.EventType)
+			}
 		}
 	}
 	return nil

pkg/repl/client_test.go

@@ -17,6 +17,7 @@ import (
 
 	"github.com/cashapp/spirit/pkg/table"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func TestMain(m *testing.M) {
@@ -547,6 +548,116 @@ func TestSetDDLNotificationChannel(t *testing.T) {
 	})
 }
 
+// TestCompositePKUpdate tests that we correctly handle
+// the case when a PRIMARY KEY is moved.
+// See: https://github.com/block/spirit/issues/417
+func TestCompositePKUpdate(t *testing.T) {
+	db, err := dbconn.New(testutils.DSN(), dbconn.NewDBConfig())
+	assert.NoError(t, err)
+	defer db.Close()
+
+	// Drop tables if they exist
+	testutils.RunSQL(t, "DROP TABLE IF EXISTS composite_pk_src, composite_pk_dst")
+
+	// Create a table with composite primary key similar to customer's message_groups table
+	testutils.RunSQL(t, `CREATE TABLE composite_pk_src (
+		organization_id BIGINT NOT NULL,
+		from_id BIGINT NOT NULL DEFAULT 0,
+		id BIGINT NOT NULL,
+		message VARCHAR(255) NOT NULL,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		PRIMARY KEY (organization_id, from_id, id),
+		UNIQUE KEY idx_id (id)
+	)`)
+
+	testutils.RunSQL(t, `CREATE TABLE composite_pk_dst (
+		organization_id BIGINT NOT NULL,
+		from_id BIGINT NOT NULL DEFAULT 0,
+		id BIGINT NOT NULL,
+		message VARCHAR(255) NOT NULL,
+		created_at DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP,
+		PRIMARY KEY (organization_id, from_id, id),
+		UNIQUE KEY idx_id (id)
+	)`)
+
+	// Insert initial test data in *both* source and destination tables
+	testutils.RunSQL(t, `INSERT INTO composite_pk_src (organization_id, from_id, id, message) VALUES 
+		(1, 100, 1, 'message 1'),
+		(1, 200, 2, 'message 2'),
+		(1, 300, 3, 'message 3'),
+		(2, 100, 4, 'message 4'),
+		(2, 200, 5, 'message 5')`)
+	testutils.RunSQL(t, `INSERT INTO composite_pk_dst SELECT * FROM composite_pk_src`)
+
+	// Set up table info
+	t1 := table.NewTableInfo(db, "test", "composite_pk_src")
+	assert.NoError(t, t1.SetInfo(t.Context()))
+	t2 := table.NewTableInfo(db, "test", "composite_pk_dst")
+	assert.NoError(t, t2.SetInfo(t.Context()))
+
+	// Create replication client
+	logger := logrus.New()
+	cfg, err := mysql2.ParseDSN(testutils.DSN())
+	assert.NoError(t, err)
+	client := NewClient(db, cfg.Addr, cfg.User, cfg.Passwd, &ClientConfig{
+		Logger:          logger,
+		Concurrency:     4,
+		TargetBatchTime: time.Second,
+		ServerID:        NewServerID(),
+	})
+
+	// Add subscription - note that keyAboveWatermark is disabled for composite PKs
+	assert.NoError(t, client.AddSubscription(t1, t2, nil))
+	assert.NoError(t, client.Run(t.Context()))
+	defer client.Close()
+
+	// Update the from_id (part of the primary key)
+	testutils.RunSQL(t, `UPDATE composite_pk_src SET from_id = 999 WHERE id IN (1, 3)`)
+	assert.NoError(t, client.BlockWait(t.Context()))
+
+	// The update should result in changes being tracked
+	// With binlog_row_image=minimal and PK updates, we expect 4 changes (2 deletes + 2 inserts)
+	deltaLen := client.GetDeltaLen()
+	require.Equal(t, 4, deltaLen, "Should have tracked 4 changes for PK update (2 deletes + 2 inserts)")
+
+	// Flush the changes
+	// This should update the destination table correctly
+	assert.NoError(t, client.Flush(t.Context()))
+
+	// Verify the data was replicated correctly
+	var count int
+
+	// Check that rows with new from_id exist in destination
+	err = db.QueryRow(`SELECT COUNT(*) FROM composite_pk_dst
+		WHERE organization_id = 1 AND from_id = 999 AND id IN (1, 3)`).Scan(&count)
+	assert.NoError(t, err)
+	assert.Equal(t, 2, count, "Rows with updated from_id should exist in destination")
+
+	// Check that rows with old from_id don't exist in destination
+	err = db.QueryRow(`SELECT COUNT(*) FROM composite_pk_dst
+		WHERE (organization_id = 1 AND from_id = 100 AND id = 1)
+		   OR (organization_id = 1 AND from_id = 300 AND id = 3)`).Scan(&count)
+	assert.NoError(t, err)
+	assert.Equal(t, 0, count, "Rows with old from_id should not exist in destination")
+
+	// Verify total row count
+	err = db.QueryRow("SELECT COUNT(*) FROM composite_pk_dst").Scan(&count)
+	assert.NoError(t, err)
+	assert.Equal(t, 5, count, "Should have all 5 rows in destination")
+
+	// Now test another PK update
+	testutils.RunSQL(t, `UPDATE composite_pk_src SET from_id = 888 WHERE id = 5`)
+	assert.NoError(t, client.BlockWait(t.Context()))
+	assert.Positive(t, client.GetDeltaLen(), "Should have tracked changes for second PK update")
+	assert.NoError(t, client.Flush(t.Context()))
+
+	// Verify the second update
+	err = db.QueryRow(`SELECT COUNT(*) FROM composite_pk_dst
+		WHERE organization_id = 2 AND from_id = 888 AND id = 5`).Scan(&count)
+	assert.NoError(t, err)
+	assert.Equal(t, 1, count, "Row with updated from_id=888 should exist in destination")
+}
+
 func TestAllChangesFlushed(t *testing.T) {
 	srcTable, dstTable := setupTestTables(t)