From 6787a8ec5878b87e4eb67e424b9401dec8b888e0 Mon Sep 17 00:00:00 2001
From: Seebs <seebs@molecula.com>
Date: Fri, 7 Apr 2023 16:23:14 -0500
Subject: [PATCH] test function bodies in parser

A naive test of a function with a body failed because we were
not actually looking for or accepting semicolons, but were generating
them on printing a thing back out. Changed to require semicolons
between statements. This logic may be wrong; for instance, if
you have a trailing semicolon on a body, the String() method won't
yield the original text exactly.
---
 sql3/parser/ast.go         |  7 +++++-
 sql3/parser/parser.go      | 13 +++++++++++
 sql3/parser/parser_test.go | 46 ++++++++++++++++++++++++++++++++++++--
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/sql3/parser/ast.go b/sql3/parser/ast.go
index 17f8738b2..befbac498 100644
--- a/sql3/parser/ast.go
+++ b/sql3/parser/ast.go
@@ -198,6 +198,8 @@ func CloneStatement(stmt Statement) Statement {
 		return stmt.Clone()
 	case *ReleaseStatement:
 		return stmt.Clone()
+	case *ReturnStatement:
+		return stmt.Clone()
 	case *RollbackStatement:
 		return stmt.Clone()
 	case *SavepointStatement:
@@ -3133,7 +3135,10 @@ func (s *CreateFunctionStatement) String() string {
 
 	buf.WriteString(" AS BEGIN")
 	for i := range s.Body {
-		fmt.Fprintf(&buf, " %s;", s.Body[i].String())
+		if i > 0 {
+			buf.WriteString(";")
+		}
+		fmt.Fprintf(&buf, " %s", s.Body[i].String())
 	}
 	buf.WriteString(" END")
 
diff --git a/sql3/parser/parser.go b/sql3/parser/parser.go
index 50533472f..7e0aaf154 100644
--- a/sql3/parser/parser.go
+++ b/sql3/parser/parser.go
@@ -1610,6 +1610,19 @@ func (p *Parser) parseCreateFunctionStatement(createPos Pos) (_ *CreateFunctionS
 		if p.peek() == END {
 			break
 		}
+		if p.peek() != SEMI {
+			return &stmt, p.errorExpected(p.pos, p.peek(), "semicolon or END")
+		}
+		// we don't have a good place to stash the semicolon's position,
+		// so just consume it.
+		_, _, _ = p.scan()
+		// Allow a trailing semicolon. Note that the trailing semicolon won't
+		// be reproduced by .String(); to fix this, we either have to modify
+		// every statement to track its own semicolon, or the body to track
+		// the positions of semicolons. Either of those might be reasonable.
+		if p.peek() == END {
+			break
+		}
 	}
 
 	if p.peek() != END {
diff --git a/sql3/parser/parser_test.go b/sql3/parser/parser_test.go
index b818a1956..165d5343b 100644
--- a/sql3/parser/parser_test.go
+++ b/sql3/parser/parser_test.go
@@ -472,7 +472,7 @@ func TestParser_ParseAlterStatement(t *testing.T) {
 
 func TestParser_ParseFunctionStatement(t *testing.T) {
 	t.Run("CreateFunction", func(t *testing.T) {
-		AssertParseStatement(t, `CREATE FUNCTION IF NOT EXISTS func (@param1 int, @param2 string) returns int as begin end`, &parser.CreateFunctionStatement{
+		AssertParseStatement(t, `CREATE FUNCTION IF NOT EXISTS func (@param1 int, @param2 string) returns int as begin return 3 end`, &parser.CreateFunctionStatement{
 			Create:      pos(0),
 			Function:    pos(7),
 			If:          pos(16),
@@ -495,8 +495,50 @@ func TestParser_ParseFunctionStatement(t *testing.T) {
 			ReturnType: &parser.Type{Name: &parser.Ident{NamePos: pos(73), Name: "int"}},
 			As:         pos(77),
 			Begin:      pos(80),
-			End:        pos(86),
+			Body: []parser.Statement{&parser.ReturnStatement{
+				Return:     pos(86),
+				ReturnExpr: &parser.IntegerLit{ValuePos: pos(93), Value: "3"},
+			}},
+			End: pos(95),
 		})
+		// This example is probably invalid, but lets us verify that the semicolons between
+		// statements are working
+		AssertParseStatement(t, `CREATE FUNCTION IF NOT EXISTS func (@param1 int, @param2 string) returns int as begin return 3; return 4 end`, &parser.CreateFunctionStatement{
+			Create:      pos(0),
+			Function:    pos(7),
+			If:          pos(16),
+			IfNot:       pos(19),
+			IfNotExists: pos(23),
+			Name:        &parser.Ident{NamePos: pos(30), Name: "func"},
+			Lparen:      pos(35),
+			Parameters: []*parser.ParameterDefinition{
+				{
+					Name: &parser.Variable{Name: "@param1", NamePos: pos(36)},
+					Type: &parser.Type{Name: &parser.Ident{NamePos: pos(44), Name: "int"}},
+				},
+				{
+					Name: &parser.Variable{Name: "@param2", NamePos: pos(49)},
+					Type: &parser.Type{Name: &parser.Ident{NamePos: pos(57), Name: "string"}},
+				},
+			},
+			Rparen:     pos(63),
+			Returns:    pos(65),
+			ReturnType: &parser.Type{Name: &parser.Ident{NamePos: pos(73), Name: "int"}},
+			As:         pos(77),
+			Begin:      pos(80),
+			Body: []parser.Statement{
+				&parser.ReturnStatement{
+					Return:     pos(86),
+					ReturnExpr: &parser.IntegerLit{ValuePos: pos(93), Value: "3"},
+				},
+				&parser.ReturnStatement{
+					Return:     pos(96),
+					ReturnExpr: &parser.IntegerLit{ValuePos: pos(103), Value: "4"},
+				},
+			},
+			End: pos(105),
+		})
+		AssertParseStatementError(t, `CREATE FUNCTION IF NOT EXISTS func (@param1 int, @param2 string) returns int as begin return 3 return 4 end`, `1:96: expected semicolon or END, found 'RETURN'`)
 		// AssertParseStatement(t, `CREATE TRIGGER IF NOT EXISTS trig BEFORE INSERT ON tbl BEGIN DELETE FROM new; END`, &parser.CreateFunctionStatement{
 		// 	Create:      pos(0),
 		// 	Function:    pos(7),