parser recognizes DROP COLUMN tokens

Author: Shlomi Noach

go/sql/parser.go

@@ -14,15 +14,18 @@ import (
 var (
 	sanitizeQuotesRegexp = regexp.MustCompile("('[^']*')")
 	renameColumnRegexp   = regexp.MustCompile(`(?i)\bchange\s+(column\s+|)([\S]+)\s+([\S]+)\s+`)
+	dropColumnRegexp     = regexp.MustCompile(`(?i)\bdrop\s+(column\s+|)([\S]+)$`)
 )
 
 type Parser struct {
 	columnRenameMap map[string]string
+	droppedColumns  map[string]bool
 }
 
 func NewParser() *Parser {
 	return &Parser{
 		columnRenameMap: make(map[string]string),
+		droppedColumns:  make(map[string]bool),
 	}
 }
 
@@ -59,10 +62,9 @@ func (this *Parser) sanitizeQuotesFromAlterStatement(alterStatement string) (str
 	return strippedStatement
 }
 
-func (this *Parser) ParseAlterStatement(alterStatement string) (err error) {
-	alterTokens, _ := this.tokenizeAlterStatement(alterStatement)
-	for _, alterToken := range alterTokens {
-		alterToken = this.sanitizeQuotesFromAlterStatement(alterToken)
+func (this *Parser) parseAlterToken(alterToken string) (err error) {
+	{
+		// rename
 		allStringSubmatch := renameColumnRegexp.FindAllStringSubmatch(alterToken, -1)
 		for _, submatch := range allStringSubmatch {
 			if unquoted, err := strconv.Unquote(submatch[2]); err == nil {
@@ -71,10 +73,28 @@ func (this *Parser) ParseAlterStatement(alterStatement string) (err error) {
 			if unquoted, err := strconv.Unquote(submatch[3]); err == nil {
 				submatch[3] = unquoted
 			}
-
 			this.columnRenameMap[submatch[2]] = submatch[3]
 		}
 	}
+	{
+		// drop
+		allStringSubmatch := dropColumnRegexp.FindAllStringSubmatch(alterToken, -1)
+		for _, submatch := range allStringSubmatch {
+			if unquoted, err := strconv.Unquote(submatch[2]); err == nil {
+				submatch[2] = unquoted
+			}
+			this.droppedColumns[submatch[2]] = true
+		}
+	}
+	return nil
+}
+
+func (this *Parser) ParseAlterStatement(alterStatement string) (err error) {
+	alterTokens, _ := this.tokenizeAlterStatement(alterStatement)
+	for _, alterToken := range alterTokens {
+		alterToken = this.sanitizeQuotesFromAlterStatement(alterToken)
+		this.parseAlterToken(alterToken)
+	}
 	return nil
 }
 

go/sql/parser_test.go

@@ -120,3 +120,42 @@ func TestSanitizeQuotesFromAlterStatement(t *testing.T) {
 		test.S(t).ExpectEquals(strippedStatement, "change column i int ''")
 	}
 }
+
+func TestParseAlterStatementDroppedColumns(t *testing.T) {
+
+	{
+		parser := NewParser()
+		statement := "drop column b"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(len(parser.droppedColumns), 1)
+		test.S(t).ExpectTrue(parser.droppedColumns["b"])
+	}
+	{
+		parser := NewParser()
+		statement := "drop column b, drop key c_idx, drop column `d`"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(len(parser.droppedColumns), 2)
+		test.S(t).ExpectTrue(parser.droppedColumns["b"])
+		test.S(t).ExpectTrue(parser.droppedColumns["d"])
+	}
+	{
+		parser := NewParser()
+		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)
+		test.S(t).ExpectEquals(len(parser.droppedColumns), 3)
+		test.S(t).ExpectTrue(parser.droppedColumns["b"])
+		test.S(t).ExpectTrue(parser.droppedColumns["d"])
+		test.S(t).ExpectTrue(parser.droppedColumns["e"])
+	}
+	{
+		parser := NewParser()
+		statement := "drop column b, drop bad statement, add column i int"
+		err := parser.ParseAlterStatement(statement)
+		test.S(t).ExpectNil(err)
+		test.S(t).ExpectEquals(len(parser.droppedColumns), 1)
+		test.S(t).ExpectTrue(parser.droppedColumns["b"])
+	}
+}