sql/ast: Resolve return values from functions (#964)

* sql/ast: Resolve return values from functions

* Fix early return with err

* Add endtoend test

Author: kylecarbs

internal/compiler/output_columns.go

@@ -229,7 +229,7 @@ func sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, error) {
 	case *ast.SelectStmt:
 		list = astutils.Search(n.FromClause, func(node ast.Node) bool {
 			switch node.(type) {
-			case *ast.RangeVar, *ast.RangeSubselect:
+			case *ast.RangeVar, *ast.RangeSubselect, *ast.FuncName:
 				return true
 			default:
 				return false
@@ -251,6 +251,20 @@ func sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, error) {
 	var tables []*Table
 	for _, item := range list.Items {
 		switch n := item.(type) {
+		case *ast.FuncName:
+			fn, err := qc.GetFunc(n)
+			if err != nil {
+				return nil, err
+			}
+			table, err := qc.GetTable(&ast.TableName{
+				Catalog: fn.ReturnType.Catalog,
+				Schema:  fn.ReturnType.Schema,
+				Name:    fn.ReturnType.Name,
+			})
+			if err != nil {
+				return nil, err
+			}
+			tables = append(tables, table)
 		case *ast.RangeSubselect:
 			cols, err := outputColumns(qc, n.Subquery)
 			if err != nil {

internal/compiler/query.go

@@ -4,6 +4,11 @@ import (
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
 )
 
+type Function struct {
+	Rel        *ast.FuncName
+	ReturnType *ast.TypeName
+}
+
 type Table struct {
 	Rel     *ast.TableName
 	Columns []*Column

internal/compiler/query_catalog.go

@@ -1,6 +1,8 @@
 package compiler
 
 import (
+	"fmt"
+
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
 	"github.com/kyleconroy/sqlc/internal/sql/catalog"
 )
@@ -71,3 +73,17 @@ func (qc QueryCatalog) GetTable(rel *ast.TableName) (*Table, error) {
 	}
 	return &Table{Rel: rel, Columns: cols}, nil
 }
+
+func (qc QueryCatalog) GetFunc(rel *ast.FuncName) (*Function, error) {
+	funcs, err := qc.catalog.ListFuncsByName(rel)
+	if err != nil {
+		return nil, err
+	}
+	if len(funcs) == 0 {
+		return nil, fmt.Errorf("function not found: %s", rel.Name)
+	}
+	return &Function{
+		Rel:        rel,
+		ReturnType: funcs[0].ReturnType,
+	}, nil
+}

internal/compiler/resolve.go

@@ -31,6 +31,23 @@ func resolveCatalogRefs(c *catalog.Catalog, rvs []*ast.RangeVar, args []paramRef
 		return defaultName
 	}
 
+	typeMap := map[string]map[string]map[string]*catalog.Column{}
+	indexTable := func(table catalog.Table) error {
+		tables = append(tables, table.Rel)
+		if defaultTable == nil {
+			defaultTable = table.Rel
+		}
+		if _, exists := typeMap[table.Rel.Schema]; !exists {
+			typeMap[table.Rel.Schema] = map[string]map[string]*catalog.Column{}
+		}
+		typeMap[table.Rel.Schema][table.Rel.Name] = map[string]*catalog.Column{}
+		for _, c := range table.Columns {
+			cc := c
+			typeMap[table.Rel.Schema][table.Rel.Name][c.Name] = cc
+		}
+		return nil
+	}
+
 	for _, rv := range rvs {
 		if rv.Relname == nil {
 			continue
@@ -39,29 +56,16 @@ func resolveCatalogRefs(c *catalog.Catalog, rvs []*ast.RangeVar, args []paramRef
 		if err != nil {
 			return nil, err
 		}
-		tables = append(tables, fqn)
-		if defaultTable == nil {
-			defaultTable = fqn
-		}
-		if rv.Alias == nil {
-			continue
-		}
-		aliasMap[*rv.Alias.Aliasname] = fqn
-	}
-
-	typeMap := map[string]map[string]map[string]*catalog.Column{}
-	for _, fqn := range tables {
 		table, err := c.GetTable(fqn)
 		if err != nil {
 			continue
 		}
-		if _, exists := typeMap[fqn.Schema]; !exists {
-			typeMap[fqn.Schema] = map[string]map[string]*catalog.Column{}
+		err = indexTable(table)
+		if err != nil {
+			return nil, err
 		}
-		typeMap[fqn.Schema][fqn.Name] = map[string]*catalog.Column{}
-		for _, c := range table.Columns {
-			cc := c
-			typeMap[fqn.Schema][fqn.Name][c.Name] = cc
+		if rv.Alias != nil {
+			aliasMap[*rv.Alias.Aliasname] = fqn
 		}
 	}
 
@@ -270,6 +274,23 @@ func resolveCatalogRefs(c *catalog.Catalog, rvs []*ast.RangeVar, args []paramRef
 				})
 			}
 
+			if fun.ReturnType == nil {
+				continue
+			}
+
+			table, err := c.GetTable(&ast.TableName{
+				Catalog: fun.ReturnType.Catalog,
+				Schema:  fun.ReturnType.Schema,
+				Name:    fun.ReturnType.Name,
+			})
+			if err != nil {
+				// The return type wasn't a table.
+				continue
+			}
+			err = indexTable(table)
+			if err != nil {
+				return nil, err
+			}
 		case *ast.ResTarget:
 			if n.Name == nil {
 				return nil, fmt.Errorf("*ast.ResTarget has nil name")

internal/endtoend/testdata/func_return/go/db.go

@@ -0,0 +1,4 @@
+/* name: GetUsers :many */
+SELECT *
+FROM users_func()
+WHERE first_name != '';

internal/endtoend/testdata/func_return/go/models.go

@@ -0,0 +1,9 @@
+CREATE TABLE users (
+    id integer,
+    first_name varchar(255) NOT NULL
+);
+
+CREATE FUNCTION users_func() RETURNS SETOF users AS $func$ BEGIN QUERY
+SELECT *
+FROM users
+END $func$ LANGUAGE plpgsql;

internal/endtoend/testdata/func_return/go/query.sql.go

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