fix(postgresql): Support selecting from a function

internal/cmd/shim.go

@@ -73,9 +73,34 @@ func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
 					Vals:    typ.Vals,
 				})
 			case *catalog.CompositeType:
+				var columns []*plugin.Column
+				for _, c := range typ.Columns {
+					l := -1
+					if c.Length != nil {
+						l = *c.Length
+					}
+					columns = append(columns, &plugin.Column{
+						Name: c.Name,
+						Type: &plugin.Identifier{
+							Catalog: c.Type.Catalog,
+							Schema:  c.Type.Schema,
+							Name:    c.Type.Name,
+						},
+						Comment:   c.Comment,
+						NotNull:   c.IsNotNull,
+						Unsigned:  c.IsUnsigned,
+						IsArray:   c.IsArray,
+						ArrayDims: int32(c.ArrayDims),
+						Length:    int32(l),
+						Table: &plugin.Identifier{
+							Name: typ.Name,
+						},
+					})
+				}
 				cts = append(cts, &plugin.CompositeType{
 					Name:    typ.Name,
 					Comment: typ.Comment,
+					Columns: columns,
 				})
 			}
 		}

internal/codegen/golang/result.go

@@ -67,40 +67,11 @@ func buildStructs(req *plugin.GenerateRequest, options *opts.Options) []Struct {
 			continue
 		}
 		for _, table := range schema.Tables {
-			var tableName string
-			if schema.Name == req.Catalog.DefaultSchema {
-				tableName = table.Rel.Name
-			} else {
-				tableName = schema.Name + "_" + table.Rel.Name
-			}
-			structName := tableName
-			if !options.EmitExactTableNames {
-				structName = inflection.Singular(inflection.SingularParams{
-					Name:       structName,
-					Exclusions: options.InflectionExcludeTableNames,
-				})
-			}
-			s := Struct{
-				Table:   &plugin.Identifier{Schema: schema.Name, Name: table.Rel.Name},
-				Name:    StructName(structName, options),
-				Comment: table.Comment,
-			}
-			for _, column := range table.Columns {
-				tags := map[string]string{}
-				if options.EmitDbTags {
-					tags["db"] = column.Name
-				}
-				if options.EmitJsonTags {
-					tags["json"] = JSONTagName(column.Name, options)
-				}
-				addExtraGoStructTags(tags, req, options, column)
-				s.Fields = append(s.Fields, Field{
-					Name:    StructName(column.Name, options),
-					Type:    goType(req, options, column),
-					Tags:    tags,
-					Comment: column.Comment,
-				})
-			}
+			s := buildStruct(req, options, schema, table.Rel.Name, table.Comment, table.Columns)
+			structs = append(structs, s)
+		}
+		for _, ty := range schema.CompositeTypes {
+			s := buildStruct(req, options, schema, ty.Name, ty.Comment, ty.Columns)
 			structs = append(structs, s)
 		}
 	}
@@ -110,6 +81,51 @@ func buildStructs(req *plugin.GenerateRequest, options *opts.Options) []Struct {
 	return structs
 }
 
+func buildStruct(
+	req *plugin.GenerateRequest,
+	options *opts.Options,
+	schema *plugin.Schema,
+	rawName string,
+	comment string,
+	columns []*plugin.Column,
+) Struct {
+	var tableName string
+	if schema.Name == req.Catalog.DefaultSchema {
+		tableName = rawName
+	} else {
+		tableName = schema.Name + "_" + rawName
+	}
+	structName := tableName
+	if !options.EmitExactTableNames {
+		structName = inflection.Singular(inflection.SingularParams{
+			Name:       structName,
+			Exclusions: options.InflectionExcludeTableNames,
+		})
+	}
+	s := Struct{
+		Table:   &plugin.Identifier{Schema: schema.Name, Name: rawName},
+		Name:    StructName(structName, options),
+		Comment: comment,
+	}
+	for _, column := range columns {
+		tags := map[string]string{}
+		if options.EmitDbTags {
+			tags["db"] = column.Name
+		}
+		if options.EmitJsonTags {
+			tags["json"] = JSONTagName(column.Name, options)
+		}
+		addExtraGoStructTags(tags, req, options, column)
+		s.Fields = append(s.Fields, Field{
+			Name:    StructName(column.Name, options),
+			Type:    goType(req, options, column),
+			Tags:    tags,
+			Comment: column.Comment,
+		})
+	}
+	return s
+}
+
 type goColumn struct {
 	id int
 	*plugin.Column

internal/compiler/output_columns.go

@@ -544,6 +544,15 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 					Schema:  fn.ReturnType.Schema,
 					Name:    fn.ReturnType.Name,
 				})
+
+				if err != nil {
+					// No table, check to see if there a type
+					tyTable, tyErr := qc.GetTableForType(fn.ReturnType)
+					if tyErr == nil {
+						table = tyTable
+						err = nil
+					}
+				}
 			}
 			if table == nil || err != nil {
 				if n.Alias != nil && len(n.Alias.Colnames.Items) > 0 {

internal/compiler/query_catalog.go

@@ -107,3 +107,23 @@ func (qc QueryCatalog) GetFunc(rel *ast.FuncName) (*Function, error) {
 		ReturnType: funcs[0].ReturnType,
 	}, nil
 }
+
+func (qc QueryCatalog) GetTableForType(rel *ast.TypeName) (*Table, error) {
+	ty, err := qc.catalog.GetCompostiteType(rel)
+	if err != nil {
+		return &Table{}, err
+	}
+
+	tblName := &ast.TableName{
+		Catalog: rel.Catalog,
+		Schema:  rel.Schema,
+		Name:    rel.Name,
+	}
+
+	var cols []*Column
+	for _, tyCol := range ty.Columns {
+		cols = append(cols, ConvertColumn(tblName, tyCol))
+	}
+
+	return &Table{Rel: tblName, Columns: cols}, nil
+}

internal/compiler/resolve.go

@@ -94,6 +94,17 @@ func (comp *Compiler) resolveCatalogRefs(qc *QueryCatalog, rvs []*ast.RangeVar,
 			continue
 		}
 
+		tyName := &ast.TypeName{
+			Catalog: embed.Table.Catalog,
+			Schema:  embed.Table.Schema,
+			Name:    embed.Table.Name,
+		}
+		ty, err := c.GetCompostiteType(tyName)
+		if err == nil {
+			embed.Table = &ast.TableName{Name: ty.Name}
+			continue
+		}
+
 		return nil, fmt.Errorf("unable to resolve table with %q: %w", embed.Orig(), err)
 	}
 

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

@@ -0,0 +1,8 @@
+-- name: SelectFoos :many
+SELECT * FROM foo_fn();
+
+-- name: SelectFooEmbed :many
+SELECT sqlc.embed(foo), 1 AS one FROM foo_fn() AS foo;
+
+-- name: SelectSingleColumn :many
+SELECT baz FROM foo_fn();