Fix backfills for tables with character-type valued primary keys (#213)

Fix #211 Add a test case and make a change to `backfill.go` to correctly quote the primary key value.
xataio · Dec 5, 2023 · 3924fd0 · 3924fd0
1 parent 788bac6
commit 3924fd0
Show file tree

Hide file tree

Showing 2 changed files with 156 additions and 70 deletions.
diff --git a/pkg/migrations/backfill.go b/pkg/migrations/backfill.go
@@ -54,7 +54,7 @@ func backfill(ctx context.Context, conn *sql.DB, table *schema.Table, cbs ...Cal
 type batcher struct {
 	table     *schema.Table
 	pkColumn  *schema.Column
-	lastPK    interface{}
+	lastPK    *string
 	batchSize int
 }
 
@@ -85,7 +85,7 @@ func (b *batcher) updateBatch(ctx context.Context, conn *sql.DB) error {
 func (b *batcher) buildQuery() string {
 	whereClause := ""
 	if b.lastPK != nil {
-		whereClause = fmt.Sprintf("WHERE %s > %v", pq.QuoteIdentifier(b.pkColumn.Name), b.lastPK)
+		whereClause = fmt.Sprintf("WHERE %s > %v", pq.QuoteIdentifier(b.pkColumn.Name), pq.QuoteLiteral(*b.lastPK))
 	}
 
 	return fmt.Sprintf(`

diff --git a/pkg/migrations/op_add_column_test.go b/pkg/migrations/op_add_column_test.go
@@ -322,88 +322,174 @@ func TestAddForeignKeyColumn(t *testing.T) {
 func TestAddColumnWithUpSql(t *testing.T) {
 	t.Parallel()
 
-	ExecuteTests(t, TestCases{{
-		name: "add column with up sql",
-		migrations: []migrations.Migration{
-			{
-				Name: "01_add_table",
-				Operations: migrations.Operations{
-					&migrations.OpCreateTable{
-						Name: "products",
-						Columns: []migrations.Column{
-							{
-								Name: "id",
-								Type: "serial",
-								Pk:   true,
-							},
-							{
-								Name:   "name",
-								Type:   "varchar(255)",
-								Unique: true,
+	ExecuteTests(t, TestCases{
+		{
+			name: "add column with up sql and a serial pk",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_add_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "products",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "serial",
+									Pk:   true,
+								},
+								{
+									Name:   "name",
+									Type:   "varchar(255)",
+									Unique: true,
+								},
 							},
 						},
 					},
 				},
-			},
-			{
-				Name: "02_add_column",
-				Operations: migrations.Operations{
-					&migrations.OpAddColumn{
-						Table: "products",
-						Up:    ptr("UPPER(name)"),
-						Column: migrations.Column{
-							Name:     "description",
-							Type:     "varchar(255)",
-							Nullable: true,
+				{
+					Name: "02_add_column",
+					Operations: migrations.Operations{
+						&migrations.OpAddColumn{
+							Table: "products",
+							Up:    ptr("UPPER(name)"),
+							Column: migrations.Column{
+								Name:     "description",
+								Type:     "varchar(255)",
+								Nullable: true,
+							},
 						},
 					},
 				},
 			},
-		},
-		afterStart: func(t *testing.T, db *sql.DB) {
-			// inserting via both the old and the new views works
-			MustInsert(t, db, "public", "01_add_table", "products", map[string]string{
-				"name": "apple",
-			})
-			MustInsert(t, db, "public", "02_add_column", "products", map[string]string{
-				"name":        "banana",
-				"description": "a yellow banana",
-			})
+			afterStart: func(t *testing.T, db *sql.DB) {
+				// inserting via both the old and the new views works
+				MustInsert(t, db, "public", "01_add_table", "products", map[string]string{
+					"name": "apple",
+				})
+				MustInsert(t, db, "public", "02_add_column", "products", map[string]string{
+					"name":        "banana",
+					"description": "a yellow banana",
+				})
 
-			res := MustSelect(t, db, "public", "02_add_column", "products")
-			assert.Equal(t, []map[string]any{
-				// the description column has been populated for the product inserted into the old view.
-				{"id": 1, "name": "apple", "description": "APPLE"},
-				// the description column for the product inserted into the new view is as inserted.
-				{"id": 2, "name": "banana", "description": "a yellow banana"},
-			}, res)
-		},
-		afterRollback: func(t *testing.T, db *sql.DB) {
-			// The trigger function has been dropped.
-			triggerFnName := migrations.TriggerFunctionName("products", "description")
-			FunctionMustNotExist(t, db, "public", triggerFnName)
+				res := MustSelect(t, db, "public", "02_add_column", "products")
+				assert.Equal(t, []map[string]any{
+					// the description column has been populated for the product inserted into the old view.
+					{"id": 1, "name": "apple", "description": "APPLE"},
+					// the description column for the product inserted into the new view is as inserted.
+					{"id": 2, "name": "banana", "description": "a yellow banana"},
+				}, res)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB) {
+				// The trigger function has been dropped.
+				triggerFnName := migrations.TriggerFunctionName("products", "description")
+				FunctionMustNotExist(t, db, "public", triggerFnName)
 
-			// The trigger has been dropped.
-			triggerName := migrations.TriggerName("products", "description")
-			TriggerMustNotExist(t, db, "public", "products", triggerName)
+				// The trigger has been dropped.
+				triggerName := migrations.TriggerName("products", "description")
+				TriggerMustNotExist(t, db, "public", "products", triggerName)
+			},
+			afterComplete: func(t *testing.T, db *sql.DB) {
+				// after rollback + restart + complete, all 'description' values are the backfilled ones.
+				res := MustSelect(t, db, "public", "02_add_column", "products")
+				assert.Equal(t, []map[string]any{
+					{"id": 1, "name": "apple", "description": "APPLE"},
+					{"id": 2, "name": "banana", "description": "BANANA"},
+				}, res)
+
+				// The trigger function has been dropped.
+				triggerFnName := migrations.TriggerFunctionName("products", "description")
+				FunctionMustNotExist(t, db, "public", triggerFnName)
+
+				// The trigger has been dropped.
+				triggerName := migrations.TriggerName("products", "description")
+				TriggerMustNotExist(t, db, "public", "products", triggerName)
+			},
 		},
-		afterComplete: func(t *testing.T, db *sql.DB) {
-			// after rollback + restart + complete, all 'description' values are the backfilled ones.
-			res := MustSelect(t, db, "public", "02_add_column", "products")
-			assert.Equal(t, []map[string]any{
-				{"id": 1, "name": "apple", "description": "APPLE"},
-				{"id": 2, "name": "banana", "description": "BANANA"},
-			}, res)
+		{
+			name: "add column with up sql and a text pk",
+			migrations: []migrations.Migration{
+				{
+					Name: "01_add_table",
+					Operations: migrations.Operations{
+						&migrations.OpCreateTable{
+							Name: "products",
+							Columns: []migrations.Column{
+								{
+									Name: "id",
+									Type: "text",
+									Pk:   true,
+								},
+								{
+									Name:   "name",
+									Type:   "varchar(255)",
+									Unique: true,
+								},
+							},
+						},
+					},
+				},
+				{
+					Name: "02_add_column",
+					Operations: migrations.Operations{
+						&migrations.OpAddColumn{
+							Table: "products",
+							Up:    ptr("UPPER(name)"),
+							Column: migrations.Column{
+								Name:     "description",
+								Type:     "varchar(255)",
+								Nullable: true,
+							},
+						},
+					},
+				},
+			},
+			afterStart: func(t *testing.T, db *sql.DB) {
+				// inserting via both the old and the new views works
+				MustInsert(t, db, "public", "01_add_table", "products", map[string]string{
+					"id":   "a",
+					"name": "apple",
+				})
+				MustInsert(t, db, "public", "02_add_column", "products", map[string]string{
+					"id":          "b",
+					"name":        "banana",
+					"description": "a yellow banana",
+				})
 
-			// The trigger function has been dropped.
-			triggerFnName := migrations.TriggerFunctionName("products", "description")
-			FunctionMustNotExist(t, db, "public", triggerFnName)
+				res := MustSelect(t, db, "public", "02_add_column", "products")
+				assert.Equal(t, []map[string]any{
+					// the description column has been populated for the product inserted into the old view.
+					{"id": "a", "name": "apple", "description": "APPLE"},
+					// the description column for the product inserted into the new view is as inserted.
+					{"id": "b", "name": "banana", "description": "a yellow banana"},
+				}, res)
+			},
+			afterRollback: func(t *testing.T, db *sql.DB) {
+				// The trigger function has been dropped.
+				triggerFnName := migrations.TriggerFunctionName("products", "description")
+				FunctionMustNotExist(t, db, "public", triggerFnName)
 
-			// The trigger has been dropped.
-			triggerName := migrations.TriggerName("products", "description")
-			TriggerMustNotExist(t, db, "public", "products", triggerName)
+				// The trigger has been dropped.
+				triggerName := migrations.TriggerName("products", "description")
+				TriggerMustNotExist(t, db, "public", "products", triggerName)
+			},
+			afterComplete: func(t *testing.T, db *sql.DB) {
+				// after rollback + restart + complete, all 'description' values are the backfilled ones.
+				res := MustSelect(t, db, "public", "02_add_column", "products")
+				assert.Equal(t, []map[string]any{
+					{"id": "a", "name": "apple", "description": "APPLE"},
+					{"id": "b", "name": "banana", "description": "BANANA"},
+				}, res)
+
+				// The trigger function has been dropped.
+				triggerFnName := migrations.TriggerFunctionName("products", "description")
+				FunctionMustNotExist(t, db, "public", triggerFnName)
+
+				// The trigger has been dropped.
+				triggerName := migrations.TriggerName("products", "description")
+				TriggerMustNotExist(t, db, "public", "products", triggerName)
+			},
 		},
-	}})
+	})
 }
 
 func TestAddNotNullColumnWithNoDefault(t *testing.T) {