Add support for --statement

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

@@ -29,7 +29,8 @@ type Resources struct {
 	DB                   *sql.DB
 	Replica              *sql.DB
 	Table                *table.TableInfo
-	Alter                string
+	Alter                string // as ALTER
+	Statement            string // as full SQL statement
 	TargetChunkTime      time.Duration
 	Threads              int
 	ReplicaMaxLag        time.Duration

pkg/check/dropadd.go

@@ -21,11 +21,10 @@ 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, "", "")
+	stmtNodes, _, err := p.Parse(r.Statement, "", "")
 	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
+		return fmt.Errorf("could not parse alter table statement: %s", r.Statement)
 	}
 	stmt := &stmtNodes[0]
 	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)

pkg/check/dropadd_test.go

@@ -4,25 +4,23 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cashapp/spirit/pkg/table"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestDropAdd(t *testing.T) {
 	r := Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "DROP b, ADD b INT",
+		Statement: "ALTER TABLE t.t1 DROP b, ADD b INT",
 	}
 	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 = "ALTER TABLE t.t1 DROP b1, ADD b2 INT"
 	err = dropAddCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err)
 
-	r.Alter = "bogus"
+	r.Statement = "bogus"
 	err = dropAddCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
 }

pkg/check/foreignkey.go

@@ -38,11 +38,10 @@ 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, "", "")
+	stmtNodes, _, err := p.Parse(r.Statement, "", "")
 	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
+		return fmt.Errorf("could not parse alter table statement: %s", r.Statement)
 	}
 	stmt := &stmtNodes[0]
 	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)

pkg/check/foreignkey_test.go

@@ -13,18 +13,17 @@ import (
 
 func TestAddForeignKey(t *testing.T) {
 	r := Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "ADD FOREIGN KEY (customer_id) REFERENCES customers (id)",
+		Statement: "ALTER TABLE t1 ADD FOREIGN KEY (customer_id) REFERENCES customers (id)",
 	}
 	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 = "ALTER TABLE t1 DROP COLUMN foo"
 	err = addForeignKeyCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err) // regular DDL
 
-	r.Alter = "bogus"
+	r.Statement = "bogus"
 	err = addForeignKeyCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // not a valid ddl
 }

pkg/check/illegalclause.go

@@ -8,12 +8,11 @@ import (
 )
 
 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 utils.AlterContainsUnsupportedClause(r.Statement)
 }

pkg/check/illegalclause_test.go

@@ -4,41 +4,30 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cashapp/spirit/pkg/table"
 	"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: "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 = "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 = "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 = "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

@@ -16,11 +16,10 @@ 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, "", "")
+	stmtNodes, _, err := p.Parse(r.Statement, "", "")
 	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
+		return fmt.Errorf("could not parse alter table statement: %s", r.Statement)
 	}
 	stmt := &stmtNodes[0]
 	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)

pkg/check/primarykey_test.go

@@ -4,29 +4,25 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cashapp/spirit/pkg/table"
 	"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: "ALTER TABLE t.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)",
+		Statement: "ALTER TABLE t.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",
+		Statement: "gibberish",
 	}
 	err = primaryKeyCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // gibberish

pkg/check/rename.go

@@ -17,11 +17,10 @@ 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, "", "")
+	stmtNodes, _, err := p.Parse(r.Statement, "", "")
 	if err != nil {
-		return fmt.Errorf("could not parse alter table statement: %s", sql)
+		return fmt.Errorf("could not parse alter table statement: %s", r.Statement)
 	}
 	stmt := &stmtNodes[0]
 	alterStmt, ok := (*stmt).(*ast.AlterTableStmt)

pkg/check/rename_test.go

@@ -4,54 +4,47 @@ import (
 	"context"
 	"testing"
 
-	"github.com/cashapp/spirit/pkg/table"
 	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 )
 
 func TestRename(t *testing.T) {
 	r := Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "RENAME TO newtablename",
+		Statement: "ALTER TABLE t.t1 RENAME TO newtablename",
 	}
 	err := renameCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
 	assert.ErrorContains(t, err, "renames are not supported")
 
 	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "RENAME COLUMN c1 TO c2",
+		Statement: "ALTER TABLE t.t1 RENAME COLUMN c1 TO c2",
 	}
 	err = renameCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
 	assert.ErrorContains(t, err, "renames are not supported")
 
 	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "CHANGE c1 c2 VARCHAR(100)",
+		Statement: "ALTER TABLE t.t1 CHANGE c1 c2 VARCHAR(100)",
 	}
 	err = renameCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err)
 	assert.ErrorContains(t, err, "renames are not supported")
 
 	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "CHANGE c1 c1 VARCHAR(100)", //nolint: dupword
+		Statement: "ALTER TABLE t.t1 CHANGE c1 c1 VARCHAR(100)", //nolint: dupword
 	}
 	err = renameCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err)
 
 	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "ADD INDEX (anothercol)",
+		Statement: "ALTER TABLE t.t1 ADD INDEX (anothercol)",
 	}
 	err = renameCheck(context.Background(), r, logrus.New())
 	assert.NoError(t, err) // safe modification
 
 	r = Resources{
-		Table: &table.TableInfo{TableName: "test"},
-		Alter: "gibberish",
+		Statement: "gibberish",
 	}
 	err = renameCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // gibberish

pkg/check/triggers.go

@@ -3,11 +3,8 @@ package check
 import (
 	"context"
 	"errors"
-	"fmt"
 	"strings"
 
-	"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,6 +18,9 @@ func init() {
 func hasTriggersCheck(ctx context.Context, r Resources, logger loggers.Advanced) error {
 	sql := `SELECT * FROM information_schema.triggers WHERE 
 	(event_object_schema=? AND event_object_table=?)`
+	if r.DB == nil {
+		return errors.New("no database connection")
+	}
 	rows, err := r.DB.QueryContext(ctx, sql, r.Table.SchemaName, r.Table.TableName)
 	if err != nil {
 		return err
@@ -36,21 +36,11 @@ func hasTriggersCheck(ctx context.Context, r Resources, logger loggers.Advanced)
 }
 
 // addTriggersCheck checks for trigger creations, which is not supported
+// Since the TiDB parser doesn't support this, we are using strings.Contains :(
 func addTriggersCheck(ctx context.Context, r Resources, logger loggers.Advanced) error {
-	isAddingTrigger := strings.Contains(strings.ToUpper(r.Alter), "CREATE TRIGGER")
+	isAddingTrigger := strings.Contains(strings.ToUpper(strings.TrimSpace(r.Statement)), "CREATE TRIGGER")
 	if isAddingTrigger {
 		return errors.New("adding triggers is not supported")
 	}
-	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]
-	_, ok := (*stmt).(*ast.AlterTableStmt)
-	if !ok {
-		return errors.New("not a valid alter table statement")
-	}
 	return nil // no problems
 }

pkg/check/triggers_test.go

@@ -16,20 +16,20 @@ import (
 
 func TestAddTriggers(t *testing.T) {
 	r := Resources{
-		Table: &table.TableInfo{TableName: "account"},
-		Alter: "CREATE TRIGGER ins_sum BEFORE INSERT ON account FOR EACH ROW SET @sum = @sum + NEW.amount;",
+		Table:     &table.TableInfo{TableName: "account"},
+		Statement: "CREATE TRIGGER ins_sum BEFORE INSERT ON account FOR EACH ROW SET @sum = @sum + NEW.amount;",
 	}
 	err := addTriggersCheck(context.Background(), r, logrus.New())
 	assert.Error(t, err) // add triggers
 	assert.ErrorContains(t, err, "adding triggers is not supported")
 
-	r.Alter = "DROP COLUMN foo"
-	err = addForeignKeyCheck(context.Background(), r, logrus.New())
+	r.Statement = "ALTER TABLE t.t1 DROP COLUMN foo"
+	err = addTriggersCheck(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
+	r.Statement = "bogus"
+	err = addTriggersCheck(context.Background(), r, logrus.New())
+	assert.NoError(t, err) // not a valid ddl, but thats ok
 }
 
 func TestHasTriggers(t *testing.T) {