allow parsing past the first delimiter for complete statement

Author: jennifersp

postgres/parser/parser/parse.go

@@ -129,15 +129,15 @@ func (p *Parser) parseOneWithDepth(depth int, sql string) (Statement, error) {
 	return stmts[0], nil
 }
 
-func (p *Parser) scanOneStmt() (sql string, tokens []sqlSymType, done bool) {
+func (p *Parser) scanOneStmt(parse func(sqlStr string, tokens []sqlSymType) error) bool {
 	var lval sqlSymType
-	tokens = p.tokBuf[:0]
+	tokens := p.tokBuf[:0]
 
 	// Scan the first token.
 	for {
 		p.scanner.scan(&lval)
 		if lval.id == 0 {
-			return "", nil, true
+			return true
 		}
 		if lval.id != ';' {
 			break
@@ -150,12 +150,21 @@ func (p *Parser) scanOneStmt() (sql string, tokens []sqlSymType, done bool) {
 	tokens = append(tokens, lval)
 	for {
 		if lval.id == ERROR {
-			return p.scanner.in[startPos:], tokens, true
+			_ = parse(p.scanner.in[startPos:], tokens)
+			return true
 		}
 		posBeforeScan := p.scanner.pos
 		p.scanner.scan(&lval)
 		if lval.id == 0 || lval.id == ';' {
-			return p.scanner.in[startPos:posBeforeScan], tokens, (lval.id == 0)
+			err := parse(p.scanner.in[startPos:posBeforeScan], tokens)
+			if lval.id == 0 || (err != nil && !strings.Contains(err.Error(), "EOF")) {
+				// done scanning all statements OR due to non EOF error
+				return true
+			} else if err == nil {
+				// done scanning single statement
+				return false
+			}
+			// continue scanning if it's EOF error
 		}
 		lval.pos -= startPos
 		tokens = append(tokens, lval)
@@ -166,19 +175,27 @@ func (p *Parser) parseWithDepth(depth int, sql string, nakedIntType *types.T) (S
 	stmts := Statements(p.stmtBuf[:0])
 	p.scanner.init(sql)
 	defer p.scanner.cleanup()
+	var err error
 	for {
-		sql, tokens, done := p.scanOneStmt()
-		stmt, err := p.parse(depth+1, sql, tokens, nakedIntType)
-		if err != nil {
-			return nil, err
-		}
-		if stmt.AST != nil {
-			stmts = append(stmts, stmt)
-		}
+		done := p.scanOneStmt(func(sqlStr string, tokens []sqlSymType) error {
+			var stmt Statement
+			stmt, err = p.parse(depth+1, sqlStr, tokens, nakedIntType)
+			if err != nil {
+				// if it's EOF syntax error, try running again
+				return err
+			}
+			if stmt.AST != nil {
+				stmts = append(stmts, stmt)
+			}
+			return nil
+		})
 		if done {
 			break
 		}
 	}
+	if err != nil {
+		return nil, err
+	}
 	return stmts, nil
 }
 
@@ -259,7 +276,9 @@ func HasMultipleStatements(sql string) bool {
 	defer p.scanner.cleanup()
 	count := 0
 	for {
-		_, _, done := p.scanOneStmt()
+		done := p.scanOneStmt(func(sqlStr string, tokens []sqlSymType) error {
+			return nil
+		})
 		if done {
 			break
 		}

postgres/parser/sem/tree/create_function.go

@@ -280,14 +280,12 @@ type BeginEndBlock struct {
 }
 
 func (node *BeginEndBlock) Format(ctx *FmtCtx) {
-	ctx.WriteString("BEGIN ATOMIC")
-	for i, s := range node.Statements {
-		if i != 0 {
-			ctx.WriteString("; ")
-		}
+	ctx.WriteString("BEGIN ATOMIC ")
+	for _, s := range node.Statements {
 		ctx.FormatNode(s)
+		ctx.WriteString("; ")
 	}
-	ctx.WriteString(" END")
+	ctx.WriteString("END")
 }
 
 var _ Statement = &Return{}

server/ast/create_function.go

@@ -112,13 +112,33 @@ func nodeCreateFunction(ctx *Context, node *tree.CreateFunction) (vitess.Stateme
 			}
 		case "sql":
 			as, ok := options[tree.OptionAs1]
-			if !ok {
-				return nil, errors.Errorf("CREATE FUNCTION definition needed for LANGUAGE SQL")
+			if ok {
+				sqlDef, sqlDefParsedStmts, err = handleLanguageSQLAs(as.Definition, params)
+				if err != nil {
+					return nil, err
+				}
+				break
 			}
-			sqlDef, sqlDefParsedStmts, err = handleLanguageSQL(as.Definition, params)
-			if err != nil {
-				return nil, err
+			sqlBody, ok := options[tree.OptionSqlBody]
+			if ok {
+				beginAtomic, ok := sqlBody.SqlBody.(*tree.BeginEndBlock)
+				if !ok {
+					return nil, errors.Errorf("Expected BEGIN ATOMIC in CREATE FUNCTION definition, got %T", sqlBody.SqlBody)
+				}
+				stmts := make([]parser.Statement, len(beginAtomic.Statements))
+				for i, s := range beginAtomic.Statements {
+					stmts[i] = parser.Statement{
+						AST: s,
+						SQL: s.String(),
+					}
+				}
+				sqlDef, sqlDefParsedStmts, err = convertSQLStmts(stmts, params)
+				if err != nil {
+					return nil, err
+				}
+				break
 			}
+			return nil, errors.Errorf("CREATE FUNCTION definition needed for LANGUAGE SQL")
 		case "c":
 			symbolOption, ok := options[tree.OptionAs2]
 			if !ok {
@@ -185,13 +205,17 @@ func createAnonymousCompositeType(fieldTypes []tree.SimpleColumnDef) *pgtypes.Do
 	return pgtypes.NewCompositeType(context.Background(), id.Null, id.NullType, typeId, attrs)
 }
 
-// handleLanguageSQL handles parsing SQL definition strings in both CREATE FUNCTION and CREATE PROCEDURE.
-func handleLanguageSQL(definition string, params []pgnodes.RoutineArg) (string, []vitess.Statement, error) {
+// handleLanguageSQLAs handles parsing SQL definition strings in both CREATE FUNCTION and CREATE PROCEDURE.
+func handleLanguageSQLAs(definition string, params []pgnodes.RoutineArg) (string, []vitess.Statement, error) {
 	stmts, err := parser.Parse(definition)
 	if err != nil {
 		return "", nil, err
 	}
 
+	return convertSQLStmts(stmts, params)
+}
+
+func convertSQLStmts(stmts []parser.Statement, params []pgnodes.RoutineArg) (string, []vitess.Statement, error) {
 	paramMap := make(map[string]*framework.ParamTypAndValue, len(params))
 	for i, param := range params {
 		tv := &framework.ParamTypAndValue{
@@ -212,14 +236,17 @@ func handleLanguageSQL(definition string, params []pgnodes.RoutineArg) (string,
 	var vitessASTs = make([]vitess.Statement, len(stmts))
 	for i, stmt := range stmts {
 		sqlDefs[i] = stmt.AST.String()
-		err = framework.ReplaceFunctionColumn(stmt.AST, paramMap)
+		err := framework.ReplaceFunctionColumn(stmt.AST, paramMap)
 		if err != nil {
 			return "", nil, err
 		}
 		// stmt.AST is updated at this point with FunctionColumn
 		vitessASTs[i], err = Convert(stmt)
+		if err != nil {
+			return "", nil, err
+		}
 	}
-	return strings.Join(sqlDefs, ";"), vitessASTs, err
+	return strings.Join(sqlDefs, ";"), vitessASTs, nil
 }
 
 // validateRoutineOptions ensures that each option is defined only once. Returns a map containing all options, or an

server/ast/create_procedure.go

@@ -106,7 +106,7 @@ func nodeCreateProcedure(ctx *Context, node *tree.CreateProcedure) (vitess.State
 			if !ok {
 				return nil, errors.Errorf("CREATE PROCEDURE definition needed for LANGUAGE SQL")
 			}
-			sqlDef, sqlDefParsedStmts, err = handleLanguageSQL(as.Definition, params)
+			sqlDef, sqlDefParsedStmts, err = handleLanguageSQLAs(as.Definition, params)
 			if err != nil {
 				return nil, err
 			}

testing/go/create_function_sql_test.go

@@ -361,5 +361,36 @@ func TestCreateFunctionsLanguageSQL(t *testing.T) {
 				},
 			},
 		},
+		{
+			Name:        "use BEGIN ATOMIC ... END in sql_body",
+			SetUpScript: []string{},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query: `CREATE FUNCTION match_default() RETURNS jsonb
+            LANGUAGE sql
+            BEGIN ATOMIC 
+				SELECT jsonb_build_object('k', 6, 'm', 2048, 'include_original', true, 'tokenizer', json_build_object('kind', 'ngram', 'token_length', 3), 'token_filters', json_build_array(json_build_object('kind', 'downcase'))) AS jsonb_build_object; 
+			END;`,
+					Expected: []sql.Row{},
+				},
+				{
+					Skip:     true, // TODO
+					Query:    `SELECT public.match_default();`,
+					Expected: []sql.Row{{`{"k": 6, "m": 2048, "tokenizer": {"kind": "ngram", "token_length": 3}, "token_filters": [{"kind": "downcase"}], "include_original": true}`}},
+				},
+				{
+					Query: `CREATE FUNCTION select1() RETURNS int
+            LANGUAGE sql
+            BEGIN ATOMIC 
+				SELECT 1; 
+			END;`,
+					Expected: []sql.Row{},
+				},
+				{
+					Query:    `SELECT select1();`,
+					Expected: []sql.Row{{1}},
+				},
+			},
+		},
 	})
 }