Implement statement hints in DMLs

ast/ast.go

@@ -577,11 +577,12 @@ type BadNode struct {
 
 // BadStatement is a BadNode for Statement.
 //
-//	{{.BadNode | sql}}
+//	{{.Hint | sqlOpt}} {{.BadNode | sql}}
 type BadStatement struct {
-	// pos = BadNode.pos
+	// pos = (Hint ?? BadNode).pos
 	// end = BadNode.end
 
+	Hint    *Hint
 	BadNode *BadNode
 }
 
@@ -592,6 +593,7 @@ type BadQueryExpr struct {
 	// pos = BadNode.pos
 	// end = BadNode.end
 
+	Hint    *Hint
 	BadNode *BadNode
 }
 
@@ -627,11 +629,12 @@ type BadDDL struct {
 
 // BadDML is a BadNode for DML.
 //
-//	{{.BadNode | sql}}
+//	{{.Hint | sqlOpt}} {{.BadNode | sql}}
 type BadDML struct {
-	// pos = BadNode.pos
+	// pos = (Hint ?? BadNode).pos
 	// end = BadNode.end
 
+	Hint    *Hint // optional
 	BadNode *BadNode
 }
 
@@ -3957,16 +3960,18 @@ type ThenReturn struct {
 
 // Insert is INSERT statement node.
 //
+//	{{.Hint | sqlOpt}}
 //	INSERT {{if .InsertOrType}}OR .InsertOrType{{end}}INTO {{.TableName | sql}} ({{.Columns | sqlJoin ","}}) {{.Input | sql}}
 //	{{.ThenReturn | sqlOpt}}
 type Insert struct {
-	// pos = Insert
+	// pos = Hint.pos || Insert
 	// end = (ThenReturn ?? Input).end
 
 	Insert token.Pos // position of "INSERT" keyword
 
 	InsertOrType InsertOrType
 
+	Hint       *Hint // optional
 	TableName  *Path
 	Columns    []*Ident
 	Input      InsertInput
@@ -4022,14 +4027,16 @@ type SubQueryInput struct {
 
 // Delete is DELETE statement.
 //
+//	{{.Hint | sqlOpt}}
 //	DELETE FROM {{.TableName | sql}} {{.As | sqlOpt}} {{.Where | sql}}
 //	{{.ThenReturn | sqlOpt}}
 type Delete struct {
-	// pos = Delete
+	// pos = Hint.pos || Delete
 	// end = (ThenReturn ?? Where).end
 
 	Delete token.Pos // position of "DELETE" keyword
 
+	Hint       *Hint // optional
 	TableName  *Path
 	As         *AsAlias // optional
 	Where      *Where
@@ -4038,15 +4045,17 @@ type Delete struct {
 
 // Update is UPDATE statement.
 //
+//	{{.Hint | sqlOpt}}
 //	UPDATE {{.TableName | sql}} {{.As | sqlOpt}}
 //	SET {{.Updates | sqlJoin ","}} {{.Where | sql}}
 //	{{.ThenReturn | sqlOpt}}
 type Update struct {
-	// pos = Update
+	// pos = Hint.pos || Update
 	// end = (ThenReturn ?? Where).end
 
 	Update token.Pos // position of "UPDATE" keyword
 
+	Hint       *Hint // optional
 	TableName  *Path
 	As         *AsAlias      // optional
 	Updates    []*UpdateItem // len(Updates) > 0

ast/sql.go

@@ -172,12 +172,12 @@ func (b *BadNode) SQL() string {
 	return sql
 }
 
-func (b *BadStatement) SQL() string { return b.BadNode.SQL() }
+func (b *BadStatement) SQL() string { return sqlOpt("", b.Hint, " ") + b.BadNode.SQL() }
 func (b *BadQueryExpr) SQL() string { return b.BadNode.SQL() }
 func (b *BadExpr) SQL() string      { return b.BadNode.SQL() }
 func (b *BadType) SQL() string      { return b.BadNode.SQL() }
 func (b *BadDDL) SQL() string       { return b.BadNode.SQL() }
-func (b *BadDML) SQL() string       { return b.BadNode.SQL() }
+func (b *BadDML) SQL() string       { return sqlOpt("", b.Hint, " ") + b.BadNode.SQL() }
 
 // ================================================================================
 //
@@ -1306,7 +1306,8 @@ func (t *ThenReturn) SQL() string {
 }
 
 func (i *Insert) SQL() string {
-	return "INSERT " +
+	return sqlOpt("", i.Hint, " ") +
+		"INSERT " +
 		strOpt(i.InsertOrType != "", "OR "+string(i.InsertOrType)+" ") +
 		"INTO " + i.TableName.SQL() + " (" +
 		sqlJoin(i.Columns, ", ") +
@@ -1335,15 +1336,17 @@ func (s *SubQueryInput) SQL() string {
 }
 
 func (d *Delete) SQL() string {
-	return "DELETE FROM " +
+	return sqlOpt("", d.Hint, " ") +
+		"DELETE FROM " +
 		d.TableName.SQL() + " " +
 		sqlOpt("", d.As, " ") +
 		d.Where.SQL() +
 		sqlOpt(" ", d.ThenReturn, "")
 }
 
 func (u *Update) SQL() string {
-	return "UPDATE " + u.TableName.SQL() + " " +
+	return sqlOpt("", u.Hint, " ") +
+		"UPDATE " + u.TableName.SQL() + " " +
 		sqlOpt("", u.As, " ") +
 		"SET " +
 		sqlJoin(u.Updates, ", ") +

parser.go

@@ -153,20 +153,38 @@ func (p *Parser) ParseDMLs() ([]ast.DML, error) {
 func (p *Parser) parseStatement() (stmt ast.Statement) {
 	l := p.Lexer.Clone()
 	defer func() {
+		// Panic on tryParseHint()
 		if r := recover(); r != nil {
 			stmt = &ast.BadStatement{BadNode: p.handleParseStatementError(r, l)}
 		}
 	}()
 
+	hint := p.tryParseHint()
+	return p.parseStatementInternal(hint)
+}
+
+func (p *Parser) parseStatementInternal(hint *ast.Hint) (stmt ast.Statement) {
+	l := p.Lexer.Clone()
+	defer func() {
+		if r := recover(); r != nil {
+			stmt = &ast.BadStatement{
+				Hint:    hint,
+				BadNode: p.handleParseStatementError(r, l),
+			}
+		}
+	}()
+
 	switch {
-	case p.Token.Kind == "SELECT" || p.Token.Kind == "@" || p.Token.Kind == "WITH" || p.Token.Kind == "(" || p.Token.Kind == "FROM":
-		return p.parseQueryStatement()
+	case p.Token.Kind == "SELECT" || p.Token.Kind == "WITH" || p.Token.Kind == "(" || p.Token.Kind == "FROM":
+		return p.parseQueryStatementInternal(hint)
+	case p.Token.IsKeywordLike("INSERT") || p.Token.IsKeywordLike("DELETE") || p.Token.IsKeywordLike("UPDATE"):
+		return p.parseDMLInternal(hint)
+	case hint != nil:
+		panic(p.errorfAtToken(&p.Token, "statement hint is only permitted before query or DML, but got: %s", p.Token.Raw))
 	case p.Token.Kind == "CREATE" || p.Token.IsKeywordLike("ALTER") || p.Token.IsKeywordLike("DROP") ||
 		p.Token.IsKeywordLike("RENAME") || p.Token.IsKeywordLike("GRANT") || p.Token.IsKeywordLike("REVOKE") ||
 		p.Token.IsKeywordLike("ANALYZE"):
 		return p.parseDDL()
-	case p.Token.IsKeywordLike("INSERT") || p.Token.IsKeywordLike("DELETE") || p.Token.IsKeywordLike("UPDATE"):
-		return p.parseDML()
 	case p.Token.IsKeywordLike("CALL"):
 		return p.parseOtherStatement()
 	}
@@ -230,15 +248,20 @@ func (p *Parser) parseQueryStatement() (stmt *ast.QueryStatement) {
 	l := p.Lexer.Clone()
 	defer func() {
 		if r := recover(); r != nil {
-			// When parsing is failed on tryParseHint or tryParseWith, the result of these methods are discarded
-			// becasue they are concrete structs and we cannot fill them with *ast.BadNode.
+			// When parsing is failed on tryParseHint, the result of these methods are discarded
+			// because they are concrete structs and we cannot fill them with *ast.BadNode.
 			stmt = &ast.QueryStatement{
 				Query: &ast.BadQueryExpr{BadNode: p.handleParseStatementError(r, l)},
 			}
 		}
 	}()
 
 	hint := p.tryParseHint()
+	return p.parseQueryStatementInternal(hint)
+}
+
+func (p *Parser) parseQueryStatementInternal(hint *ast.Hint) (stmt *ast.QueryStatement) {
+	// Can be a *ast.BadQueryExpr and won't panic
 	query := p.parseQueryExpr()
 
 	return &ast.QueryStatement{
@@ -4990,20 +5013,36 @@ func (p *Parser) parseIfExists() bool {
 func (p *Parser) parseDML() (dml ast.DML) {
 	l := p.Lexer.Clone()
 	defer func() {
+		// Panic on tryParseHint()
 		if r := recover(); r != nil {
 			dml = &ast.BadDML{BadNode: p.handleParseStatementError(r, l)}
 		}
 	}()
 
+	hint := p.tryParseHint()
+	return p.parseDMLInternal(hint)
+}
+
+func (p *Parser) parseDMLInternal(hint *ast.Hint) (dml ast.DML) {
+	l := p.Lexer.Clone()
+	defer func() {
+		if r := recover(); r != nil {
+			dml = &ast.BadDML{
+				Hint:    hint,
+				BadNode: p.handleParseStatementError(r, l),
+			}
+		}
+	}()
+
 	id := p.expect(token.TokenIdent)
 	pos := id.Pos
 	switch {
 	case id.IsKeywordLike("INSERT"):
-		return p.parseInsert(pos)
+		return p.parseInsert(pos, hint)
 	case id.IsKeywordLike("DELETE"):
-		return p.parseDelete(pos)
+		return p.parseDelete(pos, hint)
 	case id.IsKeywordLike("UPDATE"):
-		return p.parseUpdate(pos)
+		return p.parseUpdate(pos, hint)
 	}
 
 	panic(p.errorfAtToken(id, "expect pseudo keyword: INSERT, DELETE,  UPDATE but: %s", id.AsString))
@@ -5042,7 +5081,7 @@ func (p *Parser) tryParseThenReturn() *ast.ThenReturn {
 	}
 }
 
-func (p *Parser) parseInsert(pos token.Pos) *ast.Insert {
+func (p *Parser) parseInsert(pos token.Pos, hint *ast.Hint) *ast.Insert {
 	var insertOrType ast.InsertOrType
 	if p.Token.Kind == "OR" {
 		p.nextToken()
@@ -5087,6 +5126,7 @@ func (p *Parser) parseInsert(pos token.Pos) *ast.Insert {
 
 	return &ast.Insert{
 		Insert:       pos,
+		Hint:         hint,
 		InsertOrType: insertOrType,
 		TableName:    name,
 		Columns:      columns,
@@ -5151,7 +5191,7 @@ func (p *Parser) parseSubQueryInput() *ast.SubQueryInput {
 	}
 }
 
-func (p *Parser) parseDelete(pos token.Pos) *ast.Delete {
+func (p *Parser) parseDelete(pos token.Pos, hint *ast.Hint) *ast.Delete {
 	if p.Token.Kind == "FROM" {
 		p.nextToken()
 	}
@@ -5163,14 +5203,15 @@ func (p *Parser) parseDelete(pos token.Pos) *ast.Delete {
 
 	return &ast.Delete{
 		Delete:     pos,
+		Hint:       hint,
 		TableName:  name,
 		As:         as,
 		Where:      where,
 		ThenReturn: thenReturn,
 	}
 }
 
-func (p *Parser) parseUpdate(pos token.Pos) *ast.Update {
+func (p *Parser) parseUpdate(pos token.Pos, hint *ast.Hint) *ast.Update {
 	name := p.parsePath()
 	as := p.tryParseAsAlias(withOptionalAs)
 
@@ -5183,6 +5224,7 @@ func (p *Parser) parseUpdate(pos token.Pos) *ast.Update {
 
 	return &ast.Update{
 		Update:     pos,
+		Hint:       hint,
 		TableName:  name,
 		As:         as,
 		Updates:    items,

testdata/input/ddl/!bad_create_table_with_hint.sql

@@ -0,0 +1,2 @@
+@{unknown_hint=1}
+create table tbl(pk int64 primary key)

testdata/input/dml/!bad_delete.sql

@@ -0,0 +1 @@
+delete foo filter foo = 1 and bar = 2

testdata/input/dml/!bad_delete_with_bad_hint.sql

@@ -0,0 +1,2 @@
+@{invalid}
+delete foo where foo = 1 and bar = 2

testdata/input/dml/!bad_delete_with_hint.sql

@@ -0,0 +1,2 @@
+@{pdml_max_parallelism=1}
+delete foo filter foo = 1 and bar = 2

testdata/input/dml/!bad_insert_with_bad_hint.sql

@@ -0,0 +1,4 @@
+@{invalid}
+insert foo (foo, bar, baz)
+vales (1, 2, 3),
+      (4, 5, 6)

testdata/input/dml/!bad_insert_with_hint.sql

@@ -0,0 +1,4 @@
+@{pdml_max_parallelism=1}
+insert foo (foo, bar, baz)
+vales (1, 2, 3),
+      (4, 5, 6)

testdata/input/dml/!bad_update.sql

@@ -0,0 +1 @@
+update foo set invalid where foo = 1

testdata/input/dml/!bad_update_with_bad_hint.sql

@@ -0,0 +1,2 @@
+@{invalid}
+update foo set foo = bar, bar = foo, baz = DEFAULT where foo = 1

testdata/input/dml/!bad_update_with_hint.sql

@@ -0,0 +1,2 @@
+@{pdml_max_parallelism=1}
+update foo set invalid where foo = 1

testdata/input/dml/delete_with_hint.sql

@@ -0,0 +1 @@
+@{pdml_max_parallelism=1} delete foo where foo = 1 and bar = 2

testdata/input/dml/insert_with_hint.sql

@@ -0,0 +1,4 @@
+@{pdml_max_parallelism=1}
+insert into foo (foo, bar, baz)
+values (1, 2, 3),
+       (4, 5, 6)

testdata/input/dml/update_with_hint.sql

@@ -0,0 +1,2 @@
+@{pdml_max_parallelism=1}
+update foo set foo = bar, bar = foo, baz = DEFAULT where foo = 1

testdata/input/query/!bad_with_select_with_hint.sql

@@ -0,0 +1 @@
+@{hint = 1} WITH SELECT 1

testdata/input/statement/!bad_call_cancel_query_with_hint.sql

@@ -0,0 +1,2 @@
+@{unknown_hint=1}
+CALL cancel_query("12345")