feat: Allow use of table and column aliases for table functions returning unknown types (#2156)

Signed-off-by: Andrew Haines <andrew@haines.org.nz>
Co-authored-by: Kyle Conroy <kyle@conroy.org>

Author: haines

internal/compiler/output_columns.go

@@ -469,7 +469,7 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 	case *ast.SelectStmt:
 		list = astutils.Search(n.FromClause, func(node ast.Node) bool {
 			switch node.(type) {
-			case *ast.RangeVar, *ast.RangeSubselect, *ast.FuncName:
+			case *ast.RangeVar, *ast.RangeSubselect, *ast.RangeFunction:
 				return true
 			default:
 				return false
@@ -494,10 +494,20 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 	for _, item := range list.Items {
 		switch n := item.(type) {
 
-		case *ast.FuncName:
+		case *ast.RangeFunction:
 			// If the function or table can't be found, don't error out.  There
 			// are many queries that depend on functions unknown to sqlc.
-			fn, err := qc.GetFunc(n)
+			var funcCall *ast.FuncCall
+			switch f := n.Functions.Items[0].(type) {
+			case *ast.List:
+				funcCall = f.Items[0].(*ast.FuncCall)
+			case *ast.FuncCall:
+				funcCall = f
+			default:
+				return nil, fmt.Errorf("sourceTables: unsupported function call type %T", n.Functions.Items[0])
+			}
+
+			fn, err := qc.GetFunc(funcCall.Func)
 			if err != nil {
 				continue
 			}
@@ -507,7 +517,22 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 				Name:    fn.ReturnType.Name,
 			})
 			if err != nil {
-				continue
+				if n.Alias == nil || len(n.Alias.Colnames.Items) == 0 {
+					continue
+				}
+
+				table = &Table{}
+				for _, colName := range n.Alias.Colnames.Items {
+					table.Columns = append(table.Columns, &Column{
+						Name:     colName.(*ast.String).Str,
+						DataType: "any",
+					})
+				}
+			}
+			if n.Alias != nil {
+				table.Rel = &ast.TableName{
+					Name: *n.Alias.Aliasname,
+				}
 			}
 			tables = append(tables, table)
 

internal/endtoend/testdata/unnest/postgresql/pgx/v4/go/querier.go

@@ -4,3 +4,6 @@ SELECT
     unnest(@vampire_id::uuid[]) AS vampire_id
 RETURNING
     *;
+
+-- name: GetVampireIDs :many
+SELECT vampires.id::uuid FROM unnest(@vampire_id::uuid[]) AS vampires (id);

internal/endtoend/testdata/unnest/postgresql/pgx/v4/go/query.sql.go

@@ -4,3 +4,6 @@ SELECT
     unnest(@vampire_id::uuid[]) AS vampire_id
 RETURNING
     *;
+
+-- name: GetVampireIDs :many
+SELECT vampires.id::uuid FROM unnest(@vampire_id::uuid[]) AS vampires (id);

internal/endtoend/testdata/unnest/postgresql/pgx/v4/query.sql

@@ -4,3 +4,6 @@ SELECT
     unnest(@vampire_id::uuid[]) AS vampire_id
 RETURNING
     *;
+
+-- name: GetVampireIDs :many
+SELECT vampires.id::uuid FROM unnest(@vampire_id::uuid[]) AS vampires (id);