Merge pull request #375 from morgo/mtocker-add-statement-support

Add support for --statement

.gitignore

@@ -15,3 +15,6 @@ cmd/spirit/spirit
 
 # Dependency directories (remove the comment below to include it)
 # vendor/
+
+.goosehints
+.goose

README.md

@@ -123,6 +123,8 @@ Writing a new data migration tool is scary, since bugs have real consequences (d
 
 We have also tried to balance making Spirit _as fast as possible_ while still being safe to run on production systems that are running existing workloads. Sometimes this means spirit might venture into creating slow downs in application performance. If it does, please file an issue and help us make improvements.
 
+We make extensive use of the TiDB parser. If a DDL statement can not be parsed by TiDB, it will not be possible to execute it. Usually this is not a problem, but there can be [edge-cases](https://github.com/pingcap/tidb/issues/54700).
+
 ## Development
 
 See [DEVELOPMENT.md](DEVELOPMENT.md).

USAGE.md

@@ -18,6 +18,8 @@ go build
 
 The alter table command to perform. The default value is a _null alter table_, which can be useful for testing.
 
+See also: `--statement`.
+
 ### checksum
 
 - Type: Boolean
@@ -72,7 +74,7 @@ Even when force-inplace is `FALSE`, Spirit automatically detects "safe" operatio
 
 The host (and optional port) to use when connecting to MySQL.
 
-## lock-wait-timeout
+### lock-wait-timeout
 
 - Type: Duration
 - Default value: `30s`
@@ -112,6 +114,16 @@ The replication throttler only affects the copy-rows operation, and does not app
 - Temporarily disabling durability on the replica (i.e. `SET GLOBAL sync_binlog=0` and `SET GLOBAL innodb_flush_log_at_trx_commit=0`)
 - Increasing the `replica-max-lag` or disabling replica lag checking temporarily
 
+### statement
+
+- Type: String
+
+Spirit accepts either a `--table` and `--alter` argument or a `--statement` argument. When using `--statement` you can send most DDL statements to spirit, including `CREATE TABLE`, `ALTER TABLE`, `CREATE INDEX`, `RENAME TABLE` and `DROP TABLE`.
+
+The advantage of using `--statement` is you can send all schema changes directly to Spirit without having to parse statements in your automatation layer and decide which should be sent where.
+
+There are some restrictions to this. Spirit requires that the statements can be parsed by the TiDB parser, so (for example) it is not possible to send `CREATE PROCEDURE` or `CREATE TRIGGER` statements to Spirit this way.
+
 ### strict
 
 - Type: Boolean

go.mod

@@ -12,7 +12,6 @@ require (
 	github.com/siddontang/loggers v1.0.3
 	github.com/sirupsen/logrus v1.9.3
 	github.com/stretchr/testify v1.8.4
-	golang.org/x/exp v0.0.0-20231006140011-7918f672742d
 	golang.org/x/sync v0.4.0
 )
 
@@ -34,6 +33,7 @@ require (
 	go.uber.org/atomic v1.11.0 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.26.0 // indirect
+	golang.org/x/exp v0.0.0-20231006140011-7918f672742d // indirect
 	golang.org/x/sys v0.13.0 // indirect
 	golang.org/x/text v0.13.0 // indirect
 	gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect

go.sum

@@ -15,8 +15,6 @@ github.com/cznic/mathutil v0.0.0-20181122101859-297441e03548/go.mod h1:e6NPNENfs
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/go-mysql-org/go-mysql v1.8.1-0.20240805131754-ccf204bf2b2a h1:VO6kiE9ex1uNaCCgDz/q0EhTueLrr3vmxkjJpU2x6pk=
-github.com/go-mysql-org/go-mysql v1.8.1-0.20240805131754-ccf204bf2b2a/go.mod h1:+SgFgTlqjqOQoMc98n9oyUWEgn2KkOL1VmXDoq2ONOs=
 github.com/go-mysql-org/go-mysql v1.9.1 h1:W2ZKkHkoM4mmkasJCoSYfaE4RQNxXTb6VqiaMpKFrJc=
 github.com/go-mysql-org/go-mysql v1.9.1/go.mod h1:+SgFgTlqjqOQoMc98n9oyUWEgn2KkOL1VmXDoq2ONOs=
 github.com/go-sql-driver/mysql v1.7.1 h1:lUIinVbN1DY0xBg0eMOzmmtGoHwWBbvnWubQUrtU8EI=

pkg/check/check.go

@@ -8,6 +8,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/cashapp/spirit/pkg/statement"
 	"github.com/cashapp/spirit/pkg/table"
 	"github.com/siddontang/loggers"
 )
@@ -29,7 +30,7 @@ type Resources struct {
 	DB                   *sql.DB
 	Replica              *sql.DB
 	Table                *table.TableInfo
-	Alter                string
+	Statement            *statement.AbstractStatement
 	TargetChunkTime      time.Duration
 	Threads              int
 	ReplicaMaxLag        time.Duration

pkg/check/dropadd.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
 	"github.com/siddontang/loggers"
@@ -21,14 +20,7 @@ func init() {
 //   - We only allow a column name to be mentioned once across all
 //     DROP and ADD parts of the alter statement.
 func dropAddCheck(ctx context.Context, r Resources, logger loggers.Advanced) error {
-	sql := fmt.Sprintf("ALTER TABLE %s %s", r.Table.TableName, r.Alter)
-	p := parser.New()
-	stmtNodes, _, err := p.Parse(sql, "", "")
-	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
-	}
-	stmt := &stmtNodes[0]
-	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)
+	alterStmt, ok := (*r.Statement.StmtNode).(*ast.AlterTableStmt)
 	if !ok {
 		return errors.New("not a valid alter table statement")
 	}

pkg/check/dropadd_test.go

@@ -4,25 +4,21 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cashapp/spirit/pkg/table"
+	"github.com/cashapp/spirit/pkg/statement"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestDropAdd(t *testing.T) {
+	var err error
 	r := Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "DROP b, ADD b INT",
+		Statement: statement.MustNew("ALTER TABLE t1 DROP COLUMN b, ADD COLUMN b INT"),
 	}
-	err := dropAddCheck(context.Background(), r, logrus.New())
+	err = dropAddCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
 	assert.ErrorContains(t, err, "column b is mentioned 2 times in the same statement")
 
-	r.Alter = "DROP b1, ADD b2 INT"
+	r.Statement = statement.MustNew("ALTER TABLE t1 DROP b1, ADD b2 INT")
 	err = dropAddCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err)
-
-	r.Alter = "bogus"
-	err = dropAddCheck(context.Background(), r, logrus.New())
-	assert.Error(t, err)
 }

pkg/check/foreignkey.go

@@ -3,9 +3,7 @@ package check
 import (
 	"context"
 	"errors"
-	"fmt"
 
-	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
 	"github.com/siddontang/loggers"
@@ -38,14 +36,7 @@ func hasForeignKeysCheck(ctx context.Context, r Resources, logger loggers.Advanc
 }
 
 func addForeignKeyCheck(ctx context.Context, r Resources, logger loggers.Advanced) error {
-	sql := fmt.Sprintf("ALTER TABLE %s %s", r.Table.TableName, r.Alter)
-	p := parser.New()
-	stmtNodes, _, err := p.Parse(sql, "", "")
-	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
-	}
-	stmt := &stmtNodes[0]
-	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)
+	alterStmt, ok := (*r.Statement.StmtNode).(*ast.AlterTableStmt)
 	if !ok {
 		return errors.New("not a valid alter table statement")
 	}

pkg/check/foreignkey_test.go

@@ -5,28 +5,25 @@ import (
 	"database/sql"
 	"testing"
 
+	"github.com/cashapp/spirit/pkg/statement"
 	"github.com/cashapp/spirit/pkg/table"
 	"github.com/cashapp/spirit/pkg/testutils"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestAddForeignKey(t *testing.T) {
+	var err error
 	r := Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "ADD FOREIGN KEY (customer_id) REFERENCES customers (id)",
+		Statement: statement.MustNew("ALTER TABLE t1 ADD FOREIGN KEY (customer_id) REFERENCES customers (id)"),
 	}
-	err := addForeignKeyCheck(context.Background(), r, logrus.New())
+	err = addForeignKeyCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // add foreign key
 	assert.ErrorContains(t, err, "adding foreign key constraints is not supported")
 
-	r.Alter = "DROP COLUMN foo"
+	r.Statement = statement.MustNew("ALTER TABLE t1 DROP COLUMN foo")
 	err = addForeignKeyCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err) // regular DDL
-
-	r.Alter = "bogus"
-	err = addForeignKeyCheck(context.Background(), r, logrus.New())
-	assert.Error(t, err) // not a valid ddl
 }
 
 func TestHasForeignKey(t *testing.T) {
@@ -58,20 +55,22 @@ func TestHasForeignKey(t *testing.T) {
 
 	// Under this model, both customers and customer_contacts are said to have foreign keys.
 	r := Resources{
-		DB:    db,
-		Table: &table.TableInfo{SchemaName: "test", TableName: "customers"},
-		Alter: "Engine=innodb",
+		DB:        db,
+		Table:     &table.TableInfo{SchemaName: "test", TableName: "customers"},
+		Statement: statement.MustNew("ALTER TABLE customers ENGINE=innodb"),
 	}
 	err = hasForeignKeysCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // already has foreign keys.
 
 	r.Table.TableName = "customer_contacts"
+	r.Statement = statement.MustNew("ALTER TABLE customer_contacts ENGINE=innodb")
 	err = hasForeignKeysCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // already has foreign keys.
 
 	_, err = db.Exec(`drop table if exists customer_contacts`)
 	assert.NoError(t, err)
 	r.Table.TableName = "customers"
+	r.Statement = statement.MustNew("ALTER TABLE customers ENGINE=innodb")
 	err = hasForeignKeysCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err) // no longer said to have foreign keys.
 }

pkg/check/illegalclause.go

@@ -3,17 +3,15 @@ package check
 import (
 	"context"
 
-	"github.com/cashapp/spirit/pkg/utils"
 	"github.com/siddontang/loggers"
 )
 
 func init() {
-	registerCheck("illegalClause", illegalClauseCheck, ScopePreRun)
+	registerCheck("illegalClause", illegalClauseCheck, ScopePreflight)
 }
 
 // illegalClauseCheck checks for the presence of specific, unsupported
 // clauses in the ALTER statement, such as ALGORITHM= and LOCK=.
 func illegalClauseCheck(ctx context.Context, r Resources, logger loggers.Advanced) error {
-	sql := "ALTER TABLE x.x " + r.Alter
-	return utils.AlterContainsUnsupportedClause(sql)
+	return r.Statement.AlterContainsUnsupportedClause()
 }

pkg/check/illegalclause_test.go

@@ -4,41 +4,31 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cashapp/spirit/pkg/table"
+	"github.com/cashapp/spirit/pkg/statement"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestIllegalClauseCheck(t *testing.T) {
 	r := Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "ALGORITHM=INPLACE",
+		Statement: statement.MustNew("ALTER TABLE t1 ADD INDEX (b), ALGORITHM=INPLACE"),
 	}
 	err := illegalClauseCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
-	assert.ErrorContains(t, err, "ALTER contains unsupported clause")
+	assert.ErrorContains(t, err, "contains unsupported clause")
 
-	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "ALGORITHM=INPLACE, LOCK=shared",
-	}
+	r.Statement = statement.MustNew("ALTER TABLE t1  ADD c INT, ALGORITHM=INPLACE, LOCK=shared")
 	err = illegalClauseCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
-	assert.ErrorContains(t, err, "ALTER contains unsupported clause")
+	assert.ErrorContains(t, err, "contains unsupported clause")
 
-	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "lock=none",
-	}
+	r.Statement = statement.MustNew("ALTER TABLE t1  ADD c INT, lock=none")
 	err = illegalClauseCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
-	assert.ErrorContains(t, err, "ALTER contains unsupported clause")
+	assert.ErrorContains(t, err, "contains unsupported clause")
 
-	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "engine=innodb, algorithm=copy",
-	}
+	r.Statement = statement.MustNew("ALTER TABLE t1 engine=innodb, algorithm=copy")
 	err = illegalClauseCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
-	assert.ErrorContains(t, err, "ALTER contains unsupported clause")
+	assert.ErrorContains(t, err, "contains unsupported clause")
 }

pkg/check/primarykey.go

@@ -3,9 +3,7 @@ package check
 import (
 	"context"
 	"errors"
-	"fmt"
 
-	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
 	"github.com/siddontang/loggers"
@@ -16,14 +14,7 @@ func init() {
 }
 
 func primaryKeyCheck(ctx context.Context, r Resources, logger loggers.Advanced) error {
-	sql := fmt.Sprintf("ALTER TABLE %s %s", r.Table.TableName, r.Alter)
-	p := parser.New()
-	stmtNodes, _, err := p.Parse(sql, "", "")
-	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
-	}
-	stmt := &stmtNodes[0]
-	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)
+	alterStmt, ok := (*r.Statement.StmtNode).(*ast.AlterTableStmt)
 	if !ok {
 		return errors.New("not a valid alter table statement")
 	}

pkg/check/primarykey_test.go

@@ -4,30 +4,19 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cashapp/spirit/pkg/table"
+	"github.com/cashapp/spirit/pkg/statement"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestPrimaryKey(t *testing.T) {
 	r := Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "DROP PRIMARY KEY, ADD PRIMARY KEY (anothercol)",
+		Statement: statement.MustNew("ALTER TABLE t1 DROP PRIMARY KEY, ADD PRIMARY KEY (anothercol)"),
 	}
 	err := primaryKeyCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // drop primary key
 
-	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "ADD INDEX (anothercol)",
-	}
+	r.Statement = statement.MustNew("ALTER TABLE t1 ADD INDEX (anothercol)")
 	err = primaryKeyCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err) // safe modification
-
-	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "gibberish",
-	}
-	err = primaryKeyCheck(context.Background(), r, logrus.New())
-	assert.Error(t, err) // gibberish
 }

pkg/check/rename.go

@@ -3,9 +3,7 @@ package check
 import (
 	"context"
 	"errors"
-	"fmt"
 
-	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
 	"github.com/siddontang/loggers"
@@ -17,14 +15,7 @@ func init() {
 
 // renameCheck checks for any renames, which are not supported.
 func renameCheck(ctx context.Context, r Resources, logger loggers.Advanced) error {
-	sql := fmt.Sprintf("ALTER TABLE %s %s", r.Table.TableName, r.Alter)
-	p := parser.New()
-	stmtNodes, _, err := p.Parse(sql, "", "")
-	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
-	}
-	stmt := &stmtNodes[0]
-	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)
+	alterStmt, ok := (*r.Statement.StmtNode).(*ast.AlterTableStmt)
 	if !ok {
 		return errors.New("not a valid alter table statement")
 	}