fix(compiler): Support functions with OUT params (#2865)

* test: add a test case for #1654

* fix(compiler): Support functions with OUT params

* Add generated code

---------

Co-authored-by: Andrew Benton <[email protected]>

Author: kyleconroy

internal/compiler/output_columns.go

@@ -550,12 +550,15 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 			if err != nil {
 				continue
 			}
-			table, err := qc.GetTable(&ast.TableName{
-				Catalog: fn.ReturnType.Catalog,
-				Schema:  fn.ReturnType.Schema,
-				Name:    fn.ReturnType.Name,
-			})
-			if err != nil {
+			var table *Table
+			if fn.ReturnType != nil {
+				table, err = qc.GetTable(&ast.TableName{
+					Catalog: fn.ReturnType.Catalog,
+					Schema:  fn.ReturnType.Schema,
+					Name:    fn.ReturnType.Name,
+				})
+			}
+			if table == nil || err != nil {
 				if n.Alias != nil && len(n.Alias.Colnames.Items) > 0 {
 					table = &Table{}
 					for _, colName := range n.Alias.Colnames.Items {
@@ -575,12 +578,22 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 							Schema:  fn.Rel.Schema,
 							Name:    fn.Rel.Name,
 						},
-						Columns: []*Column{
+					}
+					if len(fn.Outs) > 0 {
+						for _, arg := range fn.Outs {
+							table.Columns = append(table.Columns, &Column{
+								Name:     arg.Name,
+								DataType: arg.Type.Name,
+							})
+						}
+					}
+					if fn.ReturnType != nil {
+						table.Columns = []*Column{
 							{
 								Name:     colName,
 								DataType: fn.ReturnType.Name,
 							},
-						},
+						}
 					}
 				}
 			}

internal/compiler/query.go

@@ -3,11 +3,13 @@ package compiler
 import (
 	"github.com/sqlc-dev/sqlc/internal/metadata"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
+	"github.com/sqlc-dev/sqlc/internal/sql/catalog"
 )
 
 type Function struct {
 	Rel        *ast.FuncName
 	ReturnType *ast.TypeName
+	Outs       []*catalog.Argument
 }
 
 type Table struct {

internal/compiler/query_catalog.go

@@ -103,6 +103,7 @@ func (qc QueryCatalog) GetFunc(rel *ast.FuncName) (*Function, error) {
 	}
 	return &Function{
 		Rel:        rel,
+		Outs:       funcs[0].OutArgs(),
 		ReturnType: funcs[0].ReturnType,
 	}, nil
 }

internal/endtoend/testdata/func_out_param/issue.md

@@ -0,0 +1 @@
+https://github.com/sqlc-dev/sqlc/issues/1654

internal/endtoend/testdata/func_out_param/pgx/go/db.go

@@ -0,0 +1,4 @@
+-- name: CreateAuthor :one
+SELECT * FROM add_author (
+  sqlc.arg(name), sqlc.arg(bio)
+);

internal/endtoend/testdata/func_out_param/pgx/go/models.go

@@ -0,0 +1,14 @@
+-- Example queries for sqlc
+CREATE TABLE authors (
+  id   BIGSERIAL PRIMARY KEY,
+  name text      NOT NULL,
+  bio  text
+);
+
+CREATE OR REPLACE FUNCTION add_author (name text, bio text, out id int)
+AS $$
+DECLARE
+BEGIN
+  id = 123;
+END;
+$$ LANGUAGE plpgsql;

internal/endtoend/testdata/func_out_param/pgx/go/query.sql.go

@@ -0,0 +1,10 @@
+version: "2"
+sql:
+  - engine: "postgresql"
+    schema: "schema.sql"
+    queries: "query.sql"
+    gen:
+      go:
+        package: "querytest"
+        out: "go"
+        sql_package: "pgx/v5"

internal/endtoend/testdata/func_out_param/pgx/query.sql

@@ -39,6 +39,17 @@ func (f *Function) InArgs() []*Argument {
 	return args
 }
 
+func (f *Function) OutArgs() []*Argument {
+	var args []*Argument
+	for _, a := range f.Args {
+		switch a.Mode {
+		case ast.FuncParamOut:
+			args = append(args, a)
+		}
+	}
+	return args
+}
+
 func (c *Catalog) createFunction(stmt *ast.CreateFunctionStmt) error {
 	ns := stmt.Func.Schema
 	if ns == "" {