Improvements

Signed-off-by: Matt Lord <[email protected]>
vitessio · Nov 7, 2024 · 4ceed1c · 4ceed1c
1 parent 7e97530
commit 4ceed1c
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 14 deletions.
diff --git a/go/vt/vtctl/workflow/materializer.go b/go/vt/vtctl/workflow/materializer.go
@@ -386,20 +386,31 @@ func (mz *materializer) deploySchema() error {
 				}
 
 				if removeAutoInc {
-					var replaceFunc func(columnName string)
+					var replaceFunc func(columnName string) error
 					if mz.ms.GetWorkflowOptions().ShardedAutoIncrementHandling == vtctldatapb.ShardedAutoIncrementHandling_REPLACE {
-						replaceFunc = func(columnName string) {
+						replaceFunc = func(columnName string) error {
 							mu.Lock()
 							defer mu.Unlock()
 							// At this point we've already confirmed that the table exists in the target
 							// vschema.
 							table := targetVSchema.Tables[ts.TargetTable]
 							// Don't override or redo anything that already exists.
 							if table != nil && table.AutoIncrement == nil {
-								seqTableName, _ := sqlescape.EnsureEscaped(fmt.Sprintf(autoSequenceTableFormat, ts.TargetTable))
-								if mz.ms.WorkflowOptions.GlobalKeyspace != "" {
-									seqTableKeyspace, _ := sqlescape.EnsureEscaped(mz.ms.WorkflowOptions.GlobalKeyspace)
-									seqTableName = fmt.Sprintf("%s.%s", seqTableKeyspace, seqTableName)
+								tableName, err := sqlescape.UnescapeID(ts.TargetTable)
+								if err != nil {
+									return err
+								}
+								seqTableName := fmt.Sprintf(autoSequenceTableFormat, tableName)
+								seqTableName, err = sqlescape.EnsureEscaped(seqTableName)
+								if err != nil {
+									return err
+								}
+								if mz.ms.GetWorkflowOptions().GlobalKeyspace != "" {
+									seqKeyspace, err := sqlescape.EnsureEscaped(mz.ms.WorkflowOptions.GlobalKeyspace)
+									if err != nil {
+										return err
+									}
+									seqTableName = fmt.Sprintf("%s.%s", seqKeyspace, seqTableName)
 								}
 								// Create a Vitess AutoIncrement definition -- which uses a sequence -- to
 								// replace the MySQL auto_increment definition that we removed.
@@ -409,6 +420,7 @@ func (mz *materializer) deploySchema() error {
 								}
 								updatedVSchema = true
 							}
+							return nil
 						}
 					}
 					ddl, err = stripAutoIncrement(ddl, mz.env.Parser(), replaceFunc)

diff --git a/go/vt/vtctl/workflow/materializer_env_test.go b/go/vt/vtctl/workflow/materializer_env_test.go
@@ -121,7 +121,7 @@ func newTestMaterializerEnv(t *testing.T, ctx context.Context, ms *vtctldatapb.M
 		tableName := ts.TargetTable
 		table, err := venv.Parser().TableFromStatement(ts.SourceExpression)
 		if err == nil {
-			tableName = table.Name.String()
+			tableName = sqlparser.String(table.Name)
 		}
 		var (
 			cols   []string

diff --git a/go/vt/vtctl/workflow/materializer_test.go b/go/vt/vtctl/workflow/materializer_test.go
@@ -613,12 +613,12 @@ func TestMoveTablesDDLFlag(t *testing.T) {
 // 2. REMOVE the tables' MySQL auto_increment clauses
 // 3. REPLACE the table's MySQL auto_increment clauses with Vitess sequences
 func TestShardedAutoIncHandling(t *testing.T) {
-	tableName := "t1"
+	tableName := "`t-1`"
 	tableDDL := fmt.Sprintf("create table %s (id int not null auto_increment primary key, c1 varchar(10))", tableName)
-	validateEmptyTableQuery := fmt.Sprintf("select 1 from `%s` limit 1", tableName)
+	validateEmptyTableQuery := fmt.Sprintf("select 1 from %s limit 1", tableName)
 	ms := &vtctldatapb.MaterializeSettings{
 		Workflow:       "workflow",
-		SourceKeyspace: "sourceks",
+		SourceKeyspace: "source-ks",
 		TargetKeyspace: "targetks",
 		TableSettings: []*vtctldatapb.TableMaterializeSettings{{
 			TargetTable:      tableName,
@@ -863,7 +863,7 @@ func TestShardedAutoIncHandling(t *testing.T) {
 						},
 						AutoIncrement: &vschemapb.AutoIncrement{ // AutoIncrement definition added
 							Column:   "id",
-							Sequence: fmt.Sprintf("`%s`.`%s`", ms.SourceKeyspace, fmt.Sprintf(autoSequenceTableFormat, tableName)),
+							Sequence: fmt.Sprintf("`%s`.`%s`", ms.SourceKeyspace, fmt.Sprintf(autoSequenceTableFormat, strings.ReplaceAll(tableName, "`", ""))),
 						},
 					},
 				},

diff --git a/go/vt/vtctl/workflow/utils.go b/go/vt/vtctl/workflow/utils.go
@@ -222,24 +222,30 @@ func stripTableForeignKeys(ddl string, parser *sqlparser.Parser) (string, error)
 // table definition. If an optional replace function is specified then that
 // callback will be used to e.g. replace the MySQL clause with a Vitess
 // VSchema AutoIncrement definition.
-func stripAutoIncrement(ddl string, parser *sqlparser.Parser, replace func(columnName string)) (string, error) {
+func stripAutoIncrement(ddl string, parser *sqlparser.Parser, replace func(columnName string) error) (string, error) {
 	newDDL, err := parser.ParseStrictDDL(ddl)
 	if err != nil {
 		return "", err
 	}
 
-	_ = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
+	err = sqlparser.Walk(func(node sqlparser.SQLNode) (kontinue bool, err error) {
 		switch node := node.(type) {
 		case *sqlparser.ColumnDefinition:
 			if node.Type.Options.Autoincrement {
 				node.Type.Options.Autoincrement = false
 				if replace != nil {
-					replace(sqlparser.String(node.Name))
+					if err := replace(sqlparser.String(node.Name)); err != nil {
+						return false, vterrors.Wrapf(err, "failed to replace auto_increment column %s in %q", sqlparser.String(node.Name), ddl)
+					}
+
 				}
 			}
 		}
 		return true, nil
 	}, newDDL)
+	if err != nil {
+		return "", err
+	}
 
 	return sqlparser.String(newDDL), nil
 }