support nested block statements and call statement in plpgsql (#2570)

Author: jennifersp

server/plpgsql/json.go

@@ -125,6 +125,14 @@ type plpgSQL_stmt_block struct {
 	LineNumber int32       `json:"lineno"`
 }
 
+// plpgSQL_stmt_call exists to match the expected JSON format.
+type plpgSQL_stmt_call struct {
+	LineNumber int32 `json:"lineno"`
+	Expression expr  `json:"expr"`
+	IsCall     bool  `json:"is_call"`
+	Target     datum `json:"target"`
+}
+
 // plpgSQL_stmt_case exists to match the expected JSON format.
 type plpgSQL_stmt_case struct {
 	LineNumber int32 `json:"lineno"`
@@ -274,6 +282,8 @@ type sqlstmt struct {
 // having a singular expected implementation.
 type statement struct {
 	Assignment  *plpgSQL_stmt_assign       `json:"PLpgSQL_stmt_assign"`
+	Block       *plpgSQL_stmt_block        `json:"PLpgSQL_stmt_block"`
+	Call        *plpgSQL_stmt_call         `json:"PLpgSQL_stmt_call"`
 	Case        *plpgSQL_stmt_case         `json:"PLpgSQL_stmt_case"`
 	DynExec     *plpgSQL_stmt_dynexecute   `json:"PLpgSQL_stmt_dynexecute"`
 	ExecSQL     *plpgSQL_stmt_execsql      `json:"PLpgSQL_stmt_execsql"`
@@ -310,6 +320,29 @@ func (stmt *plpgSQL_stmt_assign) Convert() (Assignment, error) {
 	}, nil
 }
 
+// Convert converts the JSON statement into its output form.
+func (stmt *plpgSQL_stmt_call) Convert() (ExecuteSQL, error) {
+	var target string
+	if !stmt.IsCall {
+		if stmt.Target.Row != nil {
+			names := make([]string, len(stmt.Target.Row.Fields))
+			for i, rowField := range stmt.Target.Row.Fields {
+				names[i] = rowField.Name
+			}
+			target = strings.Join(names, ",")
+		} else if stmt.Target.Variable != nil {
+			target = stmt.Target.Variable.RefName
+		} else {
+			return ExecuteSQL{}, errors.Errorf("unhandled datum type: %T", stmt.Target)
+		}
+	}
+	return ExecuteSQL{
+		Statement: stmt.Expression.Expression.Query,
+		Target:    target,
+	}, nil
+}
+
+// Convert converts the JSON statement into its output form.
 func (stmt *plpgSQL_stmt_case) Convert() (block Block, err error) {
 	// If the CASE statement has a main expression, start by assigning it to a variable so
 	// we can evaluate it once and only once.
@@ -430,16 +463,15 @@ func (stmt *plpgSQL_stmt_dynexecute) Convert() (DynamicExecute, error) {
 func (stmt *plpgSQL_stmt_execsql) Convert() (ExecuteSQL, error) {
 	var target string
 	if stmt.Into {
-		switch {
-		case stmt.Target.Row != nil:
+		if stmt.Target.Row != nil {
 			names := make([]string, len(stmt.Target.Row.Fields))
 			for i, rowField := range stmt.Target.Row.Fields {
 				names[i] = rowField.Name
 			}
 			target = strings.Join(names, ",")
-		case stmt.Target.Variable != nil:
+		} else if stmt.Target.Variable != nil {
 			target = stmt.Target.Variable.RefName
-		default:
+		} else {
 			return ExecuteSQL{}, errors.Errorf("unhandled datum type: %T", stmt.Target)
 		}
 	}

server/plpgsql/json_convert.go

@@ -85,6 +85,16 @@ func jsonConvertStatement(stmt statement) (Statement, error) {
 	switch {
 	case stmt.Assignment != nil:
 		return stmt.Assignment.Convert()
+	case stmt.Block != nil:
+		stmts, err := jsonConvertStatements(stmt.Block.Body)
+		if err != nil {
+			return Block{}, err
+		}
+		return Block{
+			Body: stmts,
+		}, nil
+	case stmt.Call != nil:
+		return stmt.Call.Convert()
 	case stmt.Case != nil:
 		return stmt.Case.Convert()
 	case stmt.DynExec != nil:

testing/go/create_procedure_plpgsql_test.go

@@ -441,5 +441,48 @@ $$;`,
 				},
 			},
 		},
+		{
+			Name: "use nested block statements and call statement in procedure body",
+			SetUpScript: []string{
+				`CREATE TABLE tbl (a int, b text);`,
+				`CREATE PROCEDURE add_value(IN a int, IN b text)
+            LANGUAGE plpgsql
+            AS $$
+        BEGIN
+            INSERT INTO tbl VALUES (a, b);
+        END;
+		$$;`,
+			},
+			Assertions: []ScriptTestAssertion{
+				{
+					Query: `CREATE PROCEDURE check_and_add(IN i int, IN t text)
+            LANGUAGE plpgsql
+            AS $$
+		DECLARE d text := t;
+        BEGIN
+            IF LENGTH(t) < 6 THEN
+                d = t || ' is too short';
+            END IF;
+
+            BEGIN
+                CALL add_value(i, d);
+            END;
+        END;
+        $$;`,
+				},
+				{
+					Query:    "CALL check_and_add(1, 'hi');",
+					Expected: []sql.Row{},
+				},
+				{
+					Query:    "CALL check_and_add(3, 'hellooo');",
+					Expected: []sql.Row{},
+				},
+				{
+					Query:    "SELECT * FROM tbl",
+					Expected: []sql.Row{{1, "hi is too short"}, {3, "hellooo"}},
+				},
+			},
+		},
 	})
 }