feat: Support CQ ID on the source only (#1461)

This change updates the `MigrateMessage` so that only the `_cq_id` will be a PK and it will not have a Unique constraint on it

Author: bbernays

scheduler/scheduler.go

@@ -188,7 +188,28 @@ func (s *Scheduler) Sync(ctx context.Context, client schema.ClientMeta, tables s
 	}
 
 	// send migrate messages first
-	for _, table := range tables.FlattenTables() {
+	for _, tableOriginal := range tables.FlattenTables() {
+		// var table *schema.Table
+		table := tableOriginal.Copy(nil)
+		if syncClient.deterministicCQID {
+			// No PK adjustment should occur if `_cq_id` is not present in the table
+			cqIDCol := table.Columns.Get(schema.CqIDColumn.Name)
+			if cqIDCol == nil {
+				continue
+			}
+			for i, c := range table.Columns {
+				if c.Name == schema.CqIDColumn.Name {
+					// Ensure that the cq_id column is the primary key
+					table.Columns[i].PrimaryKey = true
+					continue
+				}
+				if !c.PrimaryKey {
+					continue
+				}
+				table.Columns[i].PrimaryKey = false
+			}
+		}
+
 		res <- &message.SyncMigrateTable{
 			Table: table,
 		}

scheduler/scheduler_test.go

@@ -86,6 +86,21 @@ func testTableSuccessWithPK() *schema.Table {
 	}
 }
 
+func testTableSuccessWithCQIDPK() *schema.Table {
+	return &schema.Table{
+		Name:     "test_table_success",
+		Resolver: testResolverSuccess,
+		Columns: []schema.Column{
+			schema.CqIDColumn,
+			{
+				Name:       "test_column",
+				Type:       arrow.PrimitiveTypes.Int64,
+				PrimaryKey: true,
+			},
+		},
+	}
+}
+
 func testTableResolverPanic() *schema.Table {
 	return &schema.Table{
 		Name:     "test_table_resolver_panic",
@@ -233,6 +248,28 @@ var syncTestCases = []syncTestCase{
 		},
 		// deterministicCQID: true,
 	},
+	{
+		table: testTableSuccessWithCQIDPK(),
+		data: []scalar.Vector{
+			{
+				// This value will be validated because deterministicCQID is true
+				&scalar.UUID{Value: [16]byte{194, 83, 85, 170, 181, 44, 91, 112, 164, 224, 201, 153, 31, 90, 59, 135}, Valid: true},
+				&scalar.Int{Value: 3, Valid: true},
+			},
+		},
+		deterministicCQID: true,
+	},
+	{
+		table: testTableSuccessWithCQIDPK(),
+		data: []scalar.Vector{
+			{
+				// This value will not be validated as it will be randomly set by the scheduler
+				&scalar.UUID{},
+				&scalar.Int{Value: 3, Valid: true},
+			},
+		},
+		deterministicCQID: false,
+	},
 }
 
 func TestScheduler(t *testing.T) {
@@ -323,9 +360,10 @@ func TestScheduler_Cancellation(t *testing.T) {
 	}
 }
 
+// nolint:revive
 func testSyncTable(t *testing.T, tc syncTestCase, strategy Strategy, deterministicCQID bool) {
 	ctx := context.Background()
-	tables := []*schema.Table{}
+	var tables schema.Tables
 	if tc.table != nil {
 		tables = append(tables, tc.table)
 	}
@@ -349,11 +387,44 @@ func testSyncTable(t *testing.T, tc syncTestCase, strategy Strategy, determinist
 			record := v.Record
 			rec := tc.data[i].ToArrowRecord(record.Schema())
 			if !array.RecordEqual(rec, record) {
-				t.Fatalf("expected at i=%d: %v. got %v", i, tc.data[i], record)
+				// For records that include CqIDColumn, we can't verify equality because it is generated by the scheduler, unless deterministicCQID is true
+				onlyCqIDInequality := false
+				for col := range rec.Columns() {
+					if !deterministicCQID && rec.ColumnName(col) == schema.CqIDColumn.Name {
+						onlyCqIDInequality = true
+						continue
+					}
+					lc := rec.Column(col)
+					rc := record.Column(col)
+					if !array.Equal(lc, rc) {
+						onlyCqIDInequality = false
+					}
+				}
+				if !onlyCqIDInequality {
+					t.Fatalf("expected at i=%d: %v. got %v", i, tc.data[i], record)
+				}
 			}
 			i++
 		case *message.SyncMigrateTable:
-			// ignore
+			migratedTable := v.Table
+
+			initialTable := tables.Get(v.Table.Name)
+
+			pks := migratedTable.PrimaryKeys()
+			if deterministicCQID {
+				if len(pks) != 1 {
+					t.Fatalf("expected 1 pk. got %d", len(pks))
+				}
+				if pks[0] != schema.CqIDColumn.Name {
+					t.Fatalf("expected pk name %s. got %s", schema.CqIDColumn.Name, pks[0])
+				}
+			} else if len(pks) != len(initialTable.PrimaryKeys()) {
+				t.Fatalf("expected 0 pk. got %d", len(pks))
+			}
+
+			if len(pks) == 0 {
+				continue
+			}
 		default:
 			t.Fatalf("expected insert message. got %T", msg)
 		}