Improve support for postgresql functions that return tables:

- add IsArray field for function parameters
- support struct creation for function queries

Author: jrperritt

internal/compiler/output_columns.go

@@ -310,9 +310,22 @@ func (c *Compiler) outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, er
 			}
 			fun, err := qc.catalog.ResolveFuncCall(n)
 			if err == nil {
+				var tableCols []*catalog.Argument
+				for _, arg := range fun.Args {
+					if arg.Mode == ast.FuncParamTable {
+						tableCols = append(tableCols, arg)
+					}
+				}
+				var dt string
+				if len(tableCols) == 1 {
+					// A single column will later generate a scalar (or slice of scalar) return type.
+					dt = dataType(tableCols[0].Type)
+				} else {
+					dt = dataType(fun.ReturnType)
+				}
 				cols = append(cols, &Column{
 					Name:       name,
-					DataType:   dataType(fun.ReturnType),
+					DataType:   dt,
 					NotNull:    !fun.ReturnTypeNullable,
 					IsFuncCall: true,
 				})
@@ -526,6 +539,37 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 			if err != nil {
 				continue
 			}
+			fnsWithName, err := c.catalog.ListFuncsByName(&ast.FuncName{
+				Catalog: fn.ReturnType.Catalog,
+				Schema:  fn.ReturnType.Schema,
+				Name:    fn.ReturnType.Name,
+			})
+			// If the function was found, build a table structure to hold the columns to output.
+			if err == nil && len(fnsWithName) == 1 {
+				fnWithName := fnsWithName[0]
+				rel := &ast.TableName{
+					Catalog: fn.ReturnType.Catalog,
+					Schema:  fn.ReturnType.Schema,
+					Name:    fn.ReturnType.Name,
+				}
+				var cols []*Column
+				for _, arg := range fnWithName.Args {
+					if arg.Mode == ast.FuncParamTable {
+						col := &catalog.Column{
+							Name:    arg.Name,
+							Type:    *arg.Type,
+							IsArray: arg.IsArray,
+						}
+						convertedCol := ConvertColumn(rel, col)
+						cols = append(cols, convertedCol)
+					}
+				}
+				tables = append(tables, &Table{
+					Rel:     rel,
+					Columns: cols,
+				})
+				continue
+			}
 			table, err := qc.GetTable(&ast.TableName{
 				Catalog: fn.ReturnType.Catalog,
 				Schema:  fn.ReturnType.Schema,

internal/endtoend/testdata/func_return_table/postgresql/pgx/v5/go/db.go

@@ -0,0 +1,8 @@
+-- name: BooksByTagsViaFunc :many
+SELECT * FROM sqlc_book_schema.books_by_tags_func($1::VARCHAR[]);
+
+-- name: BooksByTagsViaFuncLimitedFields :many
+SELECT book_id, tags FROM sqlc_book_schema.books_by_tags_func($1::VARCHAR[]);
+
+-- name: CreateAuthorViaFunc :one
+SELECT * FROM create_author_func($1, $2);

internal/endtoend/testdata/func_return_table/postgresql/pgx/v5/go/models.go

@@ -0,0 +1,48 @@
+CREATE TABLE authors (
+          author_id SERIAL PRIMARY KEY,
+          name text NOT NULL DEFAULT '',
+          num integer NOT NULL DEFAULT 0
+);
+
+CREATE TABLE books (
+          book_id SERIAL PRIMARY KEY,
+          author_id integer NOT NULL REFERENCES authors(author_id),
+          isbn text NOT NULL DEFAULT '' UNIQUE,
+          book_type text NOT NULL DEFAULT 'FICTION',
+          title text NOT NULL DEFAULT '',
+          year integer NOT NULL DEFAULT 2000,
+          available timestamp with time zone NOT NULL DEFAULT 'NOW()',
+          tags varchar[] NOT NULL DEFAULT '{}'
+);
+
+-- to test query return type / schema interaction
+CREATE SCHEMA sqlc_book_schema;
+
+CREATE OR REPLACE FUNCTION sqlc_book_schema.books_by_tags_func(_tags VARCHAR[])
+RETURNS TABLE(book_id INT, title TEXT, name TEXT, isbn TEXT, tags VARCHAR[], available timestamp) AS $$
+BEGIN
+    RETURN QUERY
+    SELECT 
+        b.book_id,
+        b.title,
+        a.name,
+        b.isbn,
+        b.tags,
+        available
+    FROM books b
+    LEFT JOIN authors a ON b.author_id = a.author_id
+    WHERE b.tags && _tags
+    ORDER BY name, title ASC
+    ;
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION create_author_func(_name TEXT, _num INTEGER)
+RETURNS TABLE(id INTEGER) AS $$
+BEGIN
+    INSERT INTO authors(name, num)
+    VALUES (_name, _num)
+    RETURNING author_id into id;
+    RETURN NEXT;
+END;
+$$ LANGUAGE plpgsql;

internal/endtoend/testdata/func_return_table/postgresql/pgx/v5/go/query.sql.go

@@ -0,0 +1,13 @@
+{
+  "version": "1",
+  "packages": [
+    {
+      "path": "go",
+      "engine": "postgresql",
+      "sql_package": "pgx/v5",
+      "name": "querytest",
+      "schema": "schema.sql",
+      "queries": "query.sql"
+    }
+  ]
+}

internal/endtoend/testdata/func_return_table/postgresql/pgx/v5/query.sql

@@ -495,9 +495,10 @@ func translate(node *nodes.Node) (ast.Node, error) {
 				return nil, err
 			}
 			fp := &ast.FuncParam{
-				Name: &arg.Name,
-				Type: rel.TypeName(),
-				Mode: mode,
+				Name:    &arg.Name,
+				Type:    rel.TypeName(),
+				Mode:    mode,
+				IsArray: isArray(arg.ArgType),
 			}
 			if arg.Defexpr != nil {
 				fp.DefExpr = &ast.TODO{}

internal/endtoend/testdata/func_return_table/postgresql/pgx/v5/schema.sql

@@ -16,6 +16,7 @@ type FuncParam struct {
 	Type    *TypeName
 	DefExpr Node // Will always be &ast.TODO
 	Mode    FuncParamMode
+	IsArray bool
 }
 
 func (n *FuncParam) Pos() int {

internal/endtoend/testdata/func_return_table/postgresql/pgx/v5/sqlc.json

@@ -5,6 +5,7 @@ type FunctionParameter struct {
 	ArgType *TypeName
 	Mode    FunctionParameterMode
 	Defexpr Node
+	IsArray bool
 }
 
 func (n *FunctionParameter) Pos() int {

internal/engine/postgresql/parse.go

@@ -24,6 +24,7 @@ type Argument struct {
 	Type       *ast.TypeName
 	HasDefault bool
 	Mode       ast.FuncParamMode
+	IsArray    bool
 }
 
 func (f *Function) InArgs() []*Argument {
@@ -54,21 +55,38 @@ func (c *Catalog) createFunction(stmt *ast.CreateFunctionStmt) error {
 		ReturnType: stmt.ReturnType,
 	}
 	types := make([]*ast.TypeName, len(stmt.Params.Items))
+	var cols []*ast.ColumnDef
 	for i, item := range stmt.Params.Items {
 		arg := item.(*ast.FuncParam)
 		var name string
 		if arg.Name != nil {
 			name = *arg.Name
 		}
+		if arg.Mode == ast.FuncParamTable {
+			cols = append(cols, &ast.ColumnDef{
+				Colname:  name,
+				TypeName: arg.Type,
+				IsArray:  arg.IsArray,
+			})
+		}
 		fn.Args[i] = &Argument{
 			Name:       name,
 			Type:       arg.Type,
 			Mode:       arg.Mode,
 			HasDefault: arg.DefExpr != nil,
+			IsArray:    arg.IsArray,
 		}
 		types[i] = arg.Type
 	}
 
+	if len(cols) > 0 {
+		fn.ReturnType = &ast.TypeName{
+			Name:    stmt.Func.Name,
+			Schema:  stmt.Func.Schema,
+			Catalog: stmt.Func.Catalog,
+		}
+	}
+
 	_, idx, err := s.getFunc(stmt.Func, types)
 	if err == nil && !stmt.Replace {
 		return sqlerr.RelationExists(stmt.Func.Name)