Merge pull request #5 from openark/parse-alter-statement

Support a complete ALTER TABLE statement in --alter

go/base/context.go

@@ -77,9 +77,10 @@ func NewThrottleCheckResult(throttle bool, reason string, reasonHint ThrottleRea
 type MigrationContext struct {
 	Uuid string
 
-	DatabaseName      string
+	DatabaseName          string
-	OriginalTableName string
+	OriginalTableName     string
-	AlterStatement    string
+	AlterStatement        string
+	AlterStatementOptions string // anything following the 'ALTER TABLE [schema.]table' from AlterStatement
 
 	CountTableRows           bool
 	ConcurrentCountTableRows bool

go/cmd/gh-ost/main.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/github/gh-ost/go/base"
 	"github.com/github/gh-ost/go/logic"
+	"github.com/github/gh-ost/go/sql"
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/outbrain/golib/log"
 
@@ -172,14 +173,25 @@ func main() {
 		migrationContext.Log.SetLevel(log.ERROR)
 	}
 
+	if migrationContext.AlterStatement == "" {
+		log.Fatalf("--alter must be provided and statement must not be empty")
+	}
+	parser := sql.NewParserFromAlterStatement(migrationContext.AlterStatement)
+	migrationContext.AlterStatementOptions = parser.GetAlterStatementOptions()
+
 	if migrationContext.DatabaseName == "" {
-		migrationContext.Log.Fatalf("--database must be provided and database name must not be empty")
+		if parser.HasExplicitSchema() {
+			migrationContext.DatabaseName = parser.GetExplicitSchema()
+		} else {
+			log.Fatalf("--database must be provided and database name must not be empty, or --alter must specify database name")
+		}
 	}
 	if migrationContext.OriginalTableName == "" {
-		migrationContext.Log.Fatalf("--table must be provided and table name must not be empty")
+		if parser.HasExplicitTable() {
-	}
+			migrationContext.OriginalTableName = parser.GetExplicitTable()
-	if migrationContext.AlterStatement == "" {
+		} else {
-		migrationContext.Log.Fatalf("--alter must be provided and statement must not be empty")
+			log.Fatalf("--table must be provided and table name must not be empty, or --alter must specify table name")
+		}
 	}
 	migrationContext.Noop = !(*executeFlag)
 	if migrationContext.AllowedRunningOnMaster && migrationContext.TestOnReplica {

go/logic/applier.go

@@ -190,7 +190,7 @@ func (this *Applier) AlterGhost() error {
 	query := fmt.Sprintf(`alter /* gh-ost */ table %s.%s %s`,
 		sql.EscapeName(this.migrationContext.DatabaseName),
 		sql.EscapeName(this.migrationContext.GetGhostTableName()),
-		this.migrationContext.AlterStatement,
+		this.migrationContext.AlterStatementOptions,
 	)
 	this.migrationContext.Log.Infof("Altering ghost table %s.%s",
 		sql.EscapeName(this.migrationContext.DatabaseName),

go/logic/migrator.go

@@ -60,7 +60,7 @@ const (
 
 // Migrator is the main schema migration flow manager.
 type Migrator struct {
-	parser           *sql.Parser
+	parser           *sql.AlterTableParser
 	inspector        *Inspector
 	applier          *Applier
 	eventsStreamer   *EventsStreamer
@@ -88,7 +88,7 @@ type Migrator struct {
 func NewMigrator(context *base.MigrationContext) *Migrator {
 	migrator := &Migrator{
 		migrationContext:           context,
-		parser:                     sql.NewParser(),
+		parser:                     sql.NewAlterTableParser(),
 		ghostTableMigrated:         make(chan bool),
 		firstThrottlingCollected:   make(chan bool, 3),
 		rowCopyComplete:            make(chan error),

go/sql/parser.go

@@ -12,26 +12,54 @@ import (
 )
 
 var (
-	sanitizeQuotesRegexp = regexp.MustCompile("('[^']*')")
+	sanitizeQuotesRegexp                 = regexp.MustCompile("('[^']*')")
-	renameColumnRegexp   = regexp.MustCompile(`(?i)\bchange\s+(column\s+|)([\S]+)\s+([\S]+)\s+`)
+	renameColumnRegexp                   = regexp.MustCompile(`(?i)\bchange\s+(column\s+|)([\S]+)\s+([\S]+)\s+`)
-	dropColumnRegexp     = regexp.MustCompile(`(?i)\bdrop\s+(column\s+|)([\S]+)$`)
+	dropColumnRegexp                     = regexp.MustCompile(`(?i)\bdrop\s+(column\s+|)([\S]+)$`)
-	renameTableRegexp    = regexp.MustCompile(`(?i)\brename\s+(to|as)\s+`)
+	renameTableRegexp                    = regexp.MustCompile(`(?i)\brename\s+(to|as)\s+`)
+	alterTableExplicitSchemaTableRegexps = []*regexp.Regexp{
+		// ALTER TABLE `scm`.`tbl` something
+		regexp.MustCompile(`(?i)\balter\s+table\s+` + "`" + `([^` + "`" + `]+)` + "`" + `[.]` + "`" + `([^` + "`" + `]+)` + "`" + `\s+(.*$)`),
+		// ALTER TABLE `scm`.tbl something
+		regexp.MustCompile(`(?i)\balter\s+table\s+` + "`" + `([^` + "`" + `]+)` + "`" + `[.]([\S]+)\s+(.*$)`),
+		// ALTER TABLE scm.`tbl` something
+		regexp.MustCompile(`(?i)\balter\s+table\s+([\S]+)[.]` + "`" + `([^` + "`" + `]+)` + "`" + `\s+(.*$)`),
+		// ALTER TABLE scm.tbl something
+		regexp.MustCompile(`(?i)\balter\s+table\s+([\S]+)[.]([\S]+)\s+(.*$)`),
+	}
+	alterTableExplicitTableRegexps = []*regexp.Regexp{
+		// ALTER TABLE `tbl` something
+		regexp.MustCompile(`(?i)\balter\s+table\s+` + "`" + `([^` + "`" + `]+)` + "`" + `\s+(.*$)`),
+		// ALTER TABLE tbl something
+		regexp.MustCompile(`(?i)\balter\s+table\s+([\S]+)\s+(.*$)`),
+	}
 )
 
-type Parser struct {
+type AlterTableParser struct {
 	columnRenameMap map[string]string
 	droppedColumns  map[string]bool
 	isRenameTable   bool
+
+	alterStatementOptions string
+	alterTokens           []string
+
+	explicitSchema string
+	explicitTable  string
 }
 
-func NewParser() *Parser {
+func NewAlterTableParser() *AlterTableParser {
-	return &Parser{
+	return &AlterTableParser{
 		columnRenameMap: make(map[string]string),
 		droppedColumns:  make(map[string]bool),
 	}
 }
 
-func (this *Parser) tokenizeAlterStatement(alterStatement string) (tokens []string, err error) {
+func NewParserFromAlterStatement(alterStatement string) *AlterTableParser {
+	parser := NewAlterTableParser()
+	parser.ParseAlterStatement(alterStatement)
+	return parser
+}
+
+func (this *AlterTableParser) tokenizeAlterStatement(alterStatement string) (tokens []string, err error) {
 	terminatingQuote := rune(0)
 	f := func(c rune) bool {
 		switch {
@@ -58,13 +86,13 @@ func (this *Parser) tokenizeAlterStatement(alterStatement string) (tokens []stri
 	return tokens, nil
 }
 
-func (this *Parser) sanitizeQuotesFromAlterStatement(alterStatement string) (strippedStatement string) {
+func (this *AlterTableParser) sanitizeQuotesFromAlterStatement(alterStatement string) (strippedStatement string) {
 	strippedStatement = alterStatement
 	strippedStatement = sanitizeQuotesRegexp.ReplaceAllString(strippedStatement, "''")
 	return strippedStatement
 }
 
-func (this *Parser) parseAlterToken(alterToken string) (err error) {
+func (this *AlterTableParser) parseAlterToken(alterToken string) (err error) {
 	{
 		// rename
 		allStringSubmatch := renameColumnRegexp.FindAllStringSubmatch(alterToken, -1)
@@ -97,16 +125,34 @@ func (this *Parser) parseAlterToken(alterToken string) (err error) {
 	return nil
 }
 
-func (this *Parser) ParseAlterStatement(alterStatement string) (err error) {
+func (this *AlterTableParser) ParseAlterStatement(alterStatement string) (err error) {
-	alterTokens, _ := this.tokenizeAlterStatement(alterStatement)
+
+	this.alterStatementOptions = alterStatement
+	for _, alterTableRegexp := range alterTableExplicitSchemaTableRegexps {
+		if submatch := alterTableRegexp.FindStringSubmatch(this.alterStatementOptions); len(submatch) > 0 {
+			this.explicitSchema = submatch[1]
+			this.explicitTable = submatch[2]
+			this.alterStatementOptions = submatch[3]
+			break
+		}
+	}
+	for _, alterTableRegexp := range alterTableExplicitTableRegexps {
+		if submatch := alterTableRegexp.FindStringSubmatch(this.alterStatementOptions); len(submatch) > 0 {
+			this.explicitTable = submatch[1]
+			this.alterStatementOptions = submatch[2]
+			break
+		}
+	}
+	alterTokens, _ := this.tokenizeAlterStatement(this.alterStatementOptions)
 	for _, alterToken := range alterTokens {
 		alterToken = this.sanitizeQuotesFromAlterStatement(alterToken)
 		this.parseAlterToken(alterToken)
+		this.alterTokens = append(this.alterTokens, alterToken)
 	}
 	return nil
 }
 
-func (this *Parser) GetNonTrivialRenames() map[string]string {
+func (this *AlterTableParser) GetNonTrivialRenames() map[string]string {
 	result := make(map[string]string)
 	for column, renamed := range this.columnRenameMap {
 		if column != renamed {
@@ -116,14 +162,33 @@ func (this *Parser) GetNonTrivialRenames() map[string]string {
 	return result
 }
 
-func (this *Parser) HasNonTrivialRenames() bool {
+func (this *AlterTableParser) HasNonTrivialRenames() bool {
 	return len(this.GetNonTrivialRenames()) > 0
 }
 
-func (this *Parser) DroppedColumnsMap() map[string]bool {
+func (this *AlterTableParser) DroppedColumnsMap() map[string]bool {
 	return this.droppedColumns
 }
 
-func (this *Parser) IsRenameTable() bool {
+func (this *AlterTableParser) IsRenameTable() bool {
 	return this.isRenameTable
 }
+func (this *AlterTableParser) GetExplicitSchema() string {
+	return this.explicitSchema
+}
+
+func (this *AlterTableParser) HasExplicitSchema() bool {
+	return this.GetExplicitSchema() != ""
+}
+
+func (this *AlterTableParser) GetExplicitTable() string {
+	return this.explicitTable
+}
+
+func (this *AlterTableParser) HasExplicitTable() bool {
+	return this.GetExplicitTable() != ""
+}
+
+func (this *AlterTableParser) GetAlterStatementOptions() string {
+	return this.alterStatementOptions
+}

go/sql/parser_test.go

@@ -19,17 +19,19 @@ func init() {
 
 func TestParseAlterStatement(t *testing.T) {
 	statement := "add column t int, engine=innodb"
-	parser := NewParser()
+	parser := NewAlterTableParser()
 	err := parser.ParseAlterStatement(statement)
 	test.S(t).ExpectNil(err)
+	test.S(t).ExpectEquals(parser.alterStatementOptions, statement)
 	test.S(t).ExpectFalse(parser.HasNonTrivialRenames())
 }
 
 func TestParseAlterStatementTrivialRename(t *testing.T) {
 	statement := "add column t int, change ts ts timestamp, engine=innodb"
-	parser := NewParser()
+	parser := NewAlterTableParser()
 	err := parser.ParseAlterStatement(statement)
 	test.S(t).ExpectNil(err)
+	test.S(t).ExpectEquals(parser.alterStatementOptions, statement)
 	test.S(t).ExpectFalse(parser.HasNonTrivialRenames())
 	test.S(t).ExpectEquals(len(parser.columnRenameMap), 1)
 	test.S(t).ExpectEquals(parser.columnRenameMap["ts"], "ts")
@@ -37,9 +39,10 @@ func TestParseAlterStatementTrivialRename(t *testing.T) {
 
 func TestParseAlterStatementTrivialRenames(t *testing.T) {
 	statement := "add column t int, change ts ts timestamp, CHANGE f `f` float, engine=innodb"
-	parser := NewParser()
+	parser := NewAlterTableParser()
 	err := parser.ParseAlterStatement(statement)
 	test.S(t).ExpectNil(err)
+	test.S(t).ExpectEquals(parser.alterStatementOptions, statement)
 	test.S(t).ExpectFalse(parser.HasNonTrivialRenames())
 	test.S(t).ExpectEquals(len(parser.columnRenameMap), 2)
 	test.S(t).ExpectEquals(parser.columnRenameMap["ts"], "ts")
@@ -58,9 +61,10 @@ func TestParseAlterStatementNonTrivial(t *testing.T) {
 	}
 
 	for _, statement := range statements {
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.alterStatementOptions, statement)
 		renames := parser.GetNonTrivialRenames()
 		test.S(t).ExpectEquals(len(renames), 2)
 		test.S(t).ExpectEquals(renames["i"], "count")
@@ -69,7 +73,7 @@ func TestParseAlterStatementNonTrivial(t *testing.T) {
 }
 
 func TestTokenizeAlterStatement(t *testing.T) {
-	parser := NewParser()
+	parser := NewAlterTableParser()
 	{
 		alterStatement := "add column t int"
 		tokens, _ := parser.tokenizeAlterStatement(alterStatement)
@@ -108,7 +112,7 @@ func TestTokenizeAlterStatement(t *testing.T) {
 }
 
 func TestSanitizeQuotesFromAlterStatement(t *testing.T) {
-	parser := NewParser()
+	parser := NewAlterTableParser()
 	{
 		alterStatement := "add column e enum('a','b','c')"
 		strippedStatement := parser.sanitizeQuotesFromAlterStatement(alterStatement)
@@ -124,7 +128,7 @@ func TestSanitizeQuotesFromAlterStatement(t *testing.T) {
 func TestParseAlterStatementDroppedColumns(t *testing.T) {
 
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "drop column b"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
@@ -132,16 +136,17 @@ func TestParseAlterStatementDroppedColumns(t *testing.T) {
 		test.S(t).ExpectTrue(parser.droppedColumns["b"])
 	}
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "drop column b, drop key c_idx, drop column `d`"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.alterStatementOptions, statement)
 		test.S(t).ExpectEquals(len(parser.droppedColumns), 2)
 		test.S(t).ExpectTrue(parser.droppedColumns["b"])
 		test.S(t).ExpectTrue(parser.droppedColumns["d"])
 	}
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "drop column b, drop key c_idx, drop column `d`, drop `e`, drop primary key, drop foreign key fk_1"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
@@ -151,7 +156,7 @@ func TestParseAlterStatementDroppedColumns(t *testing.T) {
 		test.S(t).ExpectTrue(parser.droppedColumns["e"])
 	}
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "drop column b, drop bad statement, add column i int"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
@@ -163,38 +168,133 @@ func TestParseAlterStatementDroppedColumns(t *testing.T) {
 func TestParseAlterStatementRenameTable(t *testing.T) {
 
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "drop column b"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
 		test.S(t).ExpectFalse(parser.isRenameTable)
 	}
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "rename as something_else"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
 		test.S(t).ExpectTrue(parser.isRenameTable)
 	}
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "drop column b, rename as something_else"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.alterStatementOptions, statement)
 		test.S(t).ExpectTrue(parser.isRenameTable)
 	}
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "engine=innodb rename as something_else"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
 		test.S(t).ExpectTrue(parser.isRenameTable)
 	}
 	{
-		parser := NewParser()
+		parser := NewAlterTableParser()
 		statement := "rename as something_else, engine=innodb"
 		err := parser.ParseAlterStatement(statement)
 		test.S(t).ExpectNil(err)
 		test.S(t).ExpectTrue(parser.isRenameTable)
 	}
 }
+
+func TestParseAlterStatementExplicitTable(t *testing.T) {
+
+	{
+		parser := NewAlterTableParser()
+		statement := "drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "")
+		test.S(t).ExpectEquals(parser.explicitTable, "")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table tbl drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table `tbl` drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table `scm with spaces`.`tbl` drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "scm with spaces")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table `scm`.`tbl with spaces` drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "scm")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl with spaces")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table `scm`.tbl drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "scm")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table scm.`tbl` drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "scm")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table scm.tbl drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "scm")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b"}))
+	}
+	{
+		parser := NewAlterTableParser()
+		statement := "alter table scm.tbl drop column b, add index idx(i)"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(parser.explicitSchema, "scm")
+		test.S(t).ExpectEquals(parser.explicitTable, "tbl")
+		test.S(t).ExpectEquals(parser.alterStatementOptions, "drop column b, add index idx(i)")
+		test.S(t).ExpectTrue(reflect.DeepEqual(parser.alterTokens, []string{"drop column b", "add index idx(i)"}))
+	}
+}