Added support for embeding custom types

Author: wickwirew

internal/compiler/output_columns.go

@@ -545,9 +545,13 @@ func (c *Compiler) sourceTables(qc *QueryCatalog, node ast.Node) ([]*Table, erro
 					Name:    fn.ReturnType.Name,
 				})
 
-				// Failed to find table, check for type with name.
 				if err != nil {
-					table, err = qc.GetCompositeType(fn.ReturnType)
+					// 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 {

internal/compiler/query_catalog.go

@@ -108,7 +108,7 @@ func (qc QueryCatalog) GetFunc(rel *ast.FuncName) (*Function, error) {
 	}, nil
 }
 
-func (qc QueryCatalog) GetCompositeType(rel *ast.TypeName) (*Table, error) {
+func (qc QueryCatalog) GetTableForType(rel *ast.TypeName) (*Table, error) {
 	ty, err := qc.catalog.GetCompostiteType(rel)
 	if err != nil {
 		return &Table{}, err

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/go/models.go

@@ -1,5 +1,8 @@
--- name: SelectBars :many
-SELECT * FROM bar_fn();
-
 -- 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();

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

@@ -1,13 +1,8 @@
-CREATE TYPE foo AS ENUM ('bar', 'baz');
+CREATE TYPE foo AS (foo TEXT, baz INTEGER);
 
-CREATE TYPE bar AS (foo foo, baz INTEGER);
-
-CREATE FUNCTION bar_fn()
-RETURNS SETOF bar LANGUAGE SQL STABLE AS $$
-SELECT * FROM VALUES ('bar', 1);
-$$;
+CREATE TABLE bar (foo TEXT, baz INTEGER);
 
 CREATE FUNCTION foo_fn()
 RETURNS SETOF foo LANGUAGE SQL STABLE AS $$
-SELECT * FROM VALUES ('bar');
+SELECT foo, baz FROM bar;
 $$;