Simplify fetching rows with composite primary keys during backfill

go.mod

@@ -4,6 +4,7 @@ go 1.23
 
 require (
 	github.com/cloudflare/backoff v0.0.0-20240920015135-e46b80a3a7d0
+	github.com/google/uuid v1.6.0
 	github.com/lib/pq v1.10.9
 	github.com/oapi-codegen/nullable v1.1.0
 	github.com/pterm/pterm v0.12.80
@@ -41,7 +42,6 @@ require (
 	github.com/go-logr/stdr v1.2.2 // indirect
 	github.com/go-ole/go-ole v1.3.0 // indirect
 	github.com/gogo/protobuf v1.3.2 // indirect
-	github.com/google/uuid v1.6.0 // indirect
 	github.com/gookit/color v1.5.4 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect

pkg/migrations/backfill.go

@@ -134,13 +134,13 @@ func getIdentityColumns(table *schema.Table) []string {
 
 type batcher struct {
 	statementBuilder *batchStatementBuilder
-	lastValues       any
+	lastValues       []string
 }
 
 func newBatcher(table *schema.Table, batchSize int) *batcher {
 	return &batcher{
 		statementBuilder: newBatchStatementBuilder(table.Name, getIdentityColumns(table), batchSize),
-		lastValues:       nil,
+		lastValues:       make([]string, len(getIdentityColumns(table))),
 	}
 }
 
@@ -151,7 +151,11 @@ func (b *batcher) updateBatch(ctx context.Context, conn db.DB) error {
 
 		// Execute the query to update the next batch of rows and update the last PK
 		// value for the next batch
-		err := tx.QueryRowContext(ctx, query).Scan(&b.lastValues)
+		wrapper := make([]any, len(b.lastValues))
+		for i := range b.lastValues {
+			wrapper[i] = &b.lastValues[i]
+		}
+		err := tx.QueryRowContext(ctx, query).Scan(wrapper...)
 		if err != nil {
 			return err
 		}
@@ -179,51 +183,33 @@ func newBatchStatementBuilder(tableName string, identityColumnNames []string, ba
 }
 
 // buildQuery builds the query used to update the next batch of rows.
-func (sb *batchStatementBuilder) buildQuery(lastValues any) string {
+func (sb *batchStatementBuilder) buildQuery(lastValues []string) string {
 	return fmt.Sprintf("WITH batch AS (%[1]s), update AS (%[2]s) %[3]s",
 		sb.buildBatchSubQuery(lastValues),
 		sb.buildUpdateBatchSubQuery(),
 		sb.buildLastValueQuery())
 }
 
 // fetch the next batch of PK of rows to update
-func (sb *batchStatementBuilder) buildBatchSubQuery(lastValues any) string {
+func (sb *batchStatementBuilder) buildBatchSubQuery(lastValues []string) string {
 	whereClause := ""
-	if lastValues != nil {
-		conditions := make([]string, len(sb.identityColumns))
-		switch lastVals := lastValues.(type) {
-		case []int64:
-			for i, col := range sb.identityColumns {
-				conditions[i] = fmt.Sprintf("%s > %d", col, lastVals[i])
-			}
-		case []string:
-			for i, col := range sb.identityColumns {
-				conditions[i] = fmt.Sprintf("%s > %s", col, pq.QuoteLiteral(lastVals[i]))
-			}
-		case []any:
-			for i, col := range sb.identityColumns {
-				if v, ok := lastVals[i].(int); ok {
-					conditions[i] = fmt.Sprintf("%s > %d", col, v)
-				} else if v, ok := lastVals[i].(string); ok {
-					conditions[i] = fmt.Sprintf("%s > %s", col, pq.QuoteLiteral(v))
-				} else {
-					panic("unsupported type")
-				}
-			}
-		case int64:
-			conditions[0] = fmt.Sprintf("%s > %d ", sb.identityColumns[0], lastVals)
-		case string:
-			conditions[0] = fmt.Sprintf("%s > %s ", sb.identityColumns[0], pq.QuoteLiteral(lastVals))
-		default:
-			panic("unsupported type")
-		}
-		whereClause = "WHERE " + strings.Join(conditions, " AND ")
+	if len(lastValues) != 0 && lastValues[0] != "" {
+		whereClause = fmt.Sprintf("WHERE (%s) > (%s)",
+			strings.Join(sb.identityColumns, ", "), strings.Join(quoteLiteralList(lastValues), ", "))
 	}
 
 	return fmt.Sprintf("SELECT %[1]s FROM %[2]s %[3]s ORDER BY %[1]s LIMIT %[4]d FOR NO KEY UPDATE",
 		strings.Join(sb.identityColumns, ", "), sb.tableName, whereClause, sb.batchSize)
 }
 
+func quoteLiteralList(l []string) []string {
+	quoted := make([]string, len(l))
+	for i, v := range l {
+		quoted[i] = pq.QuoteLiteral(v)
+	}
+	return quoted
+}
+
 // update the rows in the batch
 func (sb *batchStatementBuilder) buildUpdateBatchSubQuery() string {
 	conditions := make([]string, len(sb.identityColumns))

pkg/migrations/backfill_test.go

@@ -13,7 +13,7 @@ func TestBatchStatementBuilder(t *testing.T) {
 		tableName       string
 		identityColumns []string
 		batchSize       int
-		lasValues       any
+		lasValues       []string
 		expected        string
 	}{
 		"single identity column no last value": {
@@ -32,29 +32,15 @@ func TestBatchStatementBuilder(t *testing.T) {
 			tableName:       "table_name",
 			identityColumns: []string{"id"},
 			batchSize:       10,
-			lasValues:       []int64{1},
-			expected:        `WITH batch AS (SELECT "id" FROM "table_name" WHERE "id" > 1 ORDER BY "id" LIMIT 10 FOR NO KEY UPDATE), update AS (UPDATE "table_name" SET "id" = "table_name"."id" FROM batch WHERE "table_name"."id" = batch."id" RETURNING "table_name"."id") SELECT LAST_VALUE("id") OVER() FROM update`,
+			lasValues:       []string{"1"},
+			expected:        `WITH batch AS (SELECT "id" FROM "table_name" WHERE ("id") > ('1') ORDER BY "id" LIMIT 10 FOR NO KEY UPDATE), update AS (UPDATE "table_name" SET "id" = "table_name"."id" FROM batch WHERE "table_name"."id" = batch."id" RETURNING "table_name"."id") SELECT LAST_VALUE("id") OVER() FROM update`,
 		},
 		"multiple identity columns with last value": {
-			tableName:       "table_name",
-			identityColumns: []string{"id", "zip"},
-			batchSize:       10,
-			lasValues:       []int64{1, 1234},
-			expected:        `WITH batch AS (SELECT "id", "zip" FROM "table_name" WHERE "id" > 1 AND "zip" > 1234 ORDER BY "id", "zip" LIMIT 10 FOR NO KEY UPDATE), update AS (UPDATE "table_name" SET "id" = "table_name"."id", "zip" = "table_name"."zip" FROM batch WHERE "table_name"."id" = batch."id" AND "table_name"."zip" = batch."zip" RETURNING "table_name"."id", "table_name"."zip") SELECT LAST_VALUE("id") OVER(), LAST_VALUE("zip") OVER() FROM update`,
-		},
-		"multiple string identity columns with last value": {
 			tableName:       "table_name",
 			identityColumns: []string{"id", "zip"},
 			batchSize:       10,
 			lasValues:       []string{"1", "1234"},
-			expected:        `WITH batch AS (SELECT "id", "zip" FROM "table_name" WHERE "id" > '1' AND "zip" > '1234' ORDER BY "id", "zip" LIMIT 10 FOR NO KEY UPDATE), update AS (UPDATE "table_name" SET "id" = "table_name"."id", "zip" = "table_name"."zip" FROM batch WHERE "table_name"."id" = batch."id" AND "table_name"."zip" = batch."zip" RETURNING "table_name"."id", "table_name"."zip") SELECT LAST_VALUE("id") OVER(), LAST_VALUE("zip") OVER() FROM update`,
-		},
-		"multiple different identity columns with last value": {
-			tableName:       "table_name",
-			identityColumns: []string{"id", "zip"},
-			batchSize:       10,
-			lasValues:       []any{1, "1234"},
-			expected:        `WITH batch AS (SELECT "id", "zip" FROM "table_name" WHERE "id" > 1 AND "zip" > '1234' ORDER BY "id", "zip" LIMIT 10 FOR NO KEY UPDATE), update AS (UPDATE "table_name" SET "id" = "table_name"."id", "zip" = "table_name"."zip" FROM batch WHERE "table_name"."id" = batch."id" AND "table_name"."zip" = batch."zip" RETURNING "table_name"."id", "table_name"."zip") SELECT LAST_VALUE("id") OVER(), LAST_VALUE("zip") OVER() FROM update`,
+			expected:        `WITH batch AS (SELECT "id", "zip" FROM "table_name" WHERE ("id", "zip") > ('1', '1234') ORDER BY "id", "zip" LIMIT 10 FOR NO KEY UPDATE), update AS (UPDATE "table_name" SET "id" = "table_name"."id", "zip" = "table_name"."zip" FROM batch WHERE "table_name"."id" = batch."id" AND "table_name"."zip" = batch."zip" RETURNING "table_name"."id", "table_name"."zip") SELECT LAST_VALUE("id") OVER(), LAST_VALUE("zip") OVER() FROM update`,
 		},
 	}
 

pkg/migrations/op_create_table_test.go

@@ -7,6 +7,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/google/uuid"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/xataio/pgroll/internal/testutils"
@@ -547,6 +548,92 @@ func TestCreateTable(t *testing.T) {
 				}, testutils.UniqueViolationErrorCode)
 			},
 		},
+		{
+			name: "create table with composite primary key and migrate it to test backfilling",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_create_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "users",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "uuid",
+									Pk:   true,
+								},
+								{
+									Name: "name",
+									Type: "text",
+									Pk:   true,
+								},
+								{
+									Name:     "city",
+									Type:     "text",
+									Nullable: true,
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_add_constraint",
+					Operations: migrations.Operations{
+						&migrations.OpCreateConstraint{
+							Name:    "nowhere_forbidden",
+							Table:   "users",
+							Columns: []string{"city"},
+							Type:    migrations.OpCreateConstraintTypeCheck,
+							Check:   ptr("city != 'nowhere'"),
+							Up: migrations.MultiColumnUpSQL{
+								"city": "'chicago'",
+							},
+							Down: migrations.MultiColumnDownSQL{
+								"city": "city",
+							},
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB, schema string) {
+				// Inserting a row into the table succeeds when the unique constraint is satisfied.
+				MustInsert(t, db, schema, "01_create_table", "users", map[string]string{
+					"id":   "00000000-0000-0000-0000-000000000001",
+					"name": "alice",
+					"city": "new york",
+				})
+				MustInsert(t, db, schema, "01_create_table", "users", map[string]string{
+					"id":   "00000000-0000-0000-0000-000000000002",
+					"name": "bob",
+					"city": "new york",
+				})
+				MustInsert(t, db, schema, "01_create_table", "users", map[string]string{
+					"id":   "00000000-0000-0000-0000-000000000003",
+					"name": "carol",
+				})
+				id1, _ := uuid.MustParse("00000000-0000-0000-0000-000000000001").MarshalText()
+				id2, _ := uuid.MustParse("00000000-0000-0000-0000-000000000002").MarshalText()
+				id3, _ := uuid.MustParse("00000000-0000-0000-0000-000000000003").MarshalText()
+				rows := MustSelect(t, db, schema, "01_create_table", "users")
+				assert.Equal(t, []map[string]any{
+					{"id": id1, "name": "alice", "city": "new york"},
+					{"id": id2, "name": "bob", "city": "new york"},
+					{"id": id3, "name": "carol", "city": nil},
+				}, rows)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB, schema string) {},
+			afterComplete: func(t *testing.T, db *sql.DB, schema string) {
+				id1, _ := uuid.MustParse("00000000-0000-0000-0000-000000000001").MarshalText()
+				id2, _ := uuid.MustParse("00000000-0000-0000-0000-000000000002").MarshalText()
+				id3, _ := uuid.MustParse("00000000-0000-0000-0000-000000000003").MarshalText()
+				rows := MustSelect(t, db, schema, "02_add_constraint", "users")
+				assert.Equal(t, []map[string]any{
+					{"id": id1, "name": "alice", "city": "chicago"},
+					{"id": id2, "name": "bob", "city": "chicago"},
+					{"id": id3, "name": "carol", "city": "chicago"},
+				}, rows)
+			},
+		},
 	})
 }