Include column types for composite type

internal/cmd/shim.go

@@ -59,6 +59,40 @@ func pluginWASM(p config.Plugin) *plugin.Codegen_WASM {
 	return nil
 }
 
+func columnSlice(cols []*catalog.Column, table *catalog.Table) []*plugin.Column {
+	var columns []*plugin.Column
+	for _, c := range cols {
+		l := -1
+		if c.Length != nil {
+			l = *c.Length
+		}
+		var tableId *plugin.Identifier = nil
+		if table != nil {
+			tableId = &plugin.Identifier{
+				Catalog: table.Rel.Catalog,
+				Schema:  table.Rel.Schema,
+				Name:    table.Rel.Name,
+			}
+		}
+		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:     tableId,
+		})
+	}
+	return columns
+}
+
 func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
 	var schemas []*plugin.Schema
 	for _, s := range c.Schemas {
@@ -76,44 +110,19 @@ func pluginCatalog(c *catalog.Catalog) *plugin.Catalog {
 				cts = append(cts, &plugin.CompositeType{
 					Name:    typ.Name,
 					Comment: typ.Comment,
+					Columns: columnSlice(typ.Columns, nil),
 				})
 			}
 		}
 		var tables []*plugin.Table
 		for _, t := range s.Tables {
-			var columns []*plugin.Column
-			for _, c := range t.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{
-						Catalog: t.Rel.Catalog,
-						Schema:  t.Rel.Schema,
-						Name:    t.Rel.Name,
-					},
-				})
-			}
 			tables = append(tables, &plugin.Table{
 				Rel: &plugin.Identifier{
 					Catalog: t.Rel.Catalog,
 					Schema:  t.Rel.Schema,
 					Name:    t.Rel.Name,
 				},
-				Columns: columns,
+				Columns: columnSlice(t.Columns, t),
 				Comment: t.Comment,
 			})
 		}

internal/codegen/golang/composite_type.go

@@ -0,0 +1,8 @@
+package golang
+
+type CompositeType struct {
+	SQLName string
+	Name    string
+	Comment string
+	Fields  []Field
+}

internal/codegen/golang/gen.go

@@ -17,13 +17,14 @@ import (
 )
 
 type tmplCtx struct {
-	Q           string
-	Package     string
-	SQLDriver   opts.SQLDriver
-	Enums       []Enum
-	Structs     []Struct
-	GoQueries   []Query
-	SqlcVersion string
+	Q              string
+	Package        string
+	SQLDriver      opts.SQLDriver
+	Enums          []Enum
+	Structs        []Struct
+	CompositeTypes []CompositeType
+	GoQueries      []Query
+	SqlcVersion    string
 
 	// TODO: Race conditions
 	SourceName string
@@ -109,12 +110,17 @@ func Generate(ctx context.Context, req *plugin.GenerateRequest) (*plugin.Generat
 	if err != nil {
 		return nil, err
 	}
+	driver := parseDriver(options.SqlPackage)
 
 	if err := opts.ValidateOpts(options); err != nil {
 		return nil, err
 	}
 
 	enums := buildEnums(req, options)
+	compositeTypes := []CompositeType{}
+	if driver.IsPGXV5() {
+		compositeTypes = buildCompositeTypes(req, options)
+	}
 	structs := buildStructs(req, options)
 	queries, err := buildQueries(req, options, structs)
 	if err != nil {
@@ -125,46 +131,49 @@ func Generate(ctx context.Context, req *plugin.GenerateRequest) (*plugin.Generat
 		enums, structs = filterUnusedStructs(enums, structs, queries)
 	}
 
-	if err := validate(options, enums, structs, queries); err != nil {
+	if err := validate(options, enums, compositeTypes, structs, queries); err != nil {
 		return nil, err
 	}
 
-	return generate(req, options, enums, structs, queries)
+	return generate(req, options, enums, compositeTypes, structs, queries)
 }
 
-func validate(options *opts.Options, enums []Enum, structs []Struct, queries []Query) error {
-	enumNames := make(map[string]struct{})
+func validate(options *opts.Options, enums []Enum, compositeTypes []CompositeType, structs []Struct, queries []Query) error {
+	usedNames := make(map[string]string)
 	for _, enum := range enums {
-		enumNames[enum.Name] = struct{}{}
-		enumNames["Null"+enum.Name] = struct{}{}
+		usedNames[enum.Name] = "enum"
+		usedNames["Null"+enum.Name] = "enum"
 	}
-	structNames := make(map[string]struct{})
 	for _, struckt := range structs {
-		if _, ok := enumNames[struckt.Name]; ok {
-			return fmt.Errorf("struct name conflicts with enum name: %s", struckt.Name)
+		if usedType, ok := usedNames[struckt.Name]; ok {
+			return fmt.Errorf("struct name conflicts with %s name: %s", usedType, struckt.Name)
+		}
+		usedNames[struckt.Name] = "struct"
+	}
+	for _, ct := range compositeTypes {
+		if usedType, ok := usedNames[ct.Name]; ok {
+			return fmt.Errorf("composite type name conflicts with %s name: %s", usedType, ct.Name)
 		}
-		structNames[struckt.Name] = struct{}{}
+		usedNames[ct.Name] = "composite type"
 	}
 	if !options.EmitExportedQueries {
 		return nil
 	}
 	for _, query := range queries {
-		if _, ok := enumNames[query.ConstantName]; ok {
-			return fmt.Errorf("query constant name conflicts with enum name: %s", query.ConstantName)
-		}
-		if _, ok := structNames[query.ConstantName]; ok {
-			return fmt.Errorf("query constant name conflicts with struct name: %s", query.ConstantName)
+		if usedType, ok := usedNames[query.ConstantName]; ok {
+			return fmt.Errorf("query constant name conflicts with %s name: %s", usedType, query.ConstantName)
 		}
 	}
 	return nil
 }
 
-func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum, structs []Struct, queries []Query) (*plugin.GenerateResponse, error) {
+func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum, compositeTypes []CompositeType, structs []Struct, queries []Query) (*plugin.GenerateResponse, error) {
 	i := &importer{
-		Options: options,
-		Queries: queries,
-		Enums:   enums,
-		Structs: structs,
+		Options:        options,
+		Queries:        queries,
+		Enums:          enums,
+		CompositeTypes: compositeTypes,
+		Structs:        structs,
 	}
 
 	tctx := tmplCtx{
@@ -183,6 +192,7 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 		Q:                         "`",
 		Package:                   options.Package,
 		Enums:                     enums,
+		CompositeTypes:            compositeTypes,
 		Structs:                   structs,
 		SqlcVersion:               req.SqlcVersion,
 		BuildTags:                 options.BuildTags,

internal/codegen/golang/imports.go

@@ -58,10 +58,11 @@ func mergeImports(imps ...fileImports) [][]ImportSpec {
 }
 
 type importer struct {
-	Options *opts.Options
-	Queries []Query
-	Enums   []Enum
-	Structs []Struct
+	Options        *opts.Options
+	Queries        []Query
+	Enums          []Enum
+	CompositeTypes []CompositeType
+	Structs        []Struct
 }
 
 func (i *importer) usesType(typ string) bool {
@@ -72,6 +73,13 @@ func (i *importer) usesType(typ string) bool {
 			}
 		}
 	}
+	for _, ct := range i.CompositeTypes {
+		for _, f := range ct.Fields {
+			if hasPrefixIgnoringSliceAndPointerPrefix(f.Type, typ) {
+				return true
+			}
+		}
+	}
 	return false
 }
 
@@ -132,6 +140,12 @@ func (i *importer) dbImports() fileImports {
 	case opts.SQLDriverPGXV5:
 		pkg = append(pkg, ImportSpec{Path: "github.com/jackc/pgx/v5/pgconn"})
 		pkg = append(pkg, ImportSpec{Path: "github.com/jackc/pgx/v5"})
+		if len(i.CompositeTypes) > 0 {
+			pkg = append(pkg, ImportSpec{Path: "github.com/jackc/pgx/v5/pgtype"})
+			if !i.Options.EmitMethodsWithDbArgument {
+				pkg = append(pkg, ImportSpec{Path: "fmt"})
+			}
+		}
 	default:
 		std = append(std, ImportSpec{Path: "database/sql"})
 		if i.Options.EmitPreparedQueries {

internal/codegen/golang/opts/enum.go

@@ -48,6 +48,10 @@ func (d SQLDriver) IsPGX() bool {
 	return d == SQLDriverPGXV4 || d == SQLDriverPGXV5
 }
 
+func (d SQLDriver) IsPGXV5() bool {
+	return d == SQLDriverPGXV5
+}
+
 func (d SQLDriver) IsGoSQLDriverMySQL() bool {
 	return d == SQLDriverGoSQLDriverMySQL
 }

internal/codegen/golang/postgresql_type.go

@@ -587,13 +587,26 @@ func postgresType(req *plugin.GenerateRequest, options *opts.Options, col *plugi
 
 			for _, ct := range schema.CompositeTypes {
 				if rel.Name == ct.Name && rel.Schema == schema.Name {
-					if notNull {
-						return "string"
+					if !driver.IsPGXV5() {
+						if notNull {
+							return "string"
+						}
+						if emitPointersForNull {
+							return "*string"
+						}
+						return "sql.NullString"
 					}
-					if emitPointersForNull {
-						return "*string"
+
+					var ctName string
+					if schema.Name == req.Catalog.DefaultSchema {
+						ctName = StructName(ct.Name, options)
+					} else {
+						ctName = StructName(schema.Name+"_"+ct.Name, options)
+					}
+					if notNull {
+						return ctName
 					}
-					return "sql.NullString"
+					return "*" + ctName
 				}
 			}
 		}

internal/codegen/golang/result.go

@@ -60,6 +60,82 @@ func buildEnums(req *plugin.GenerateRequest, options *opts.Options) []Enum {
 	return enums
 }
 
+func topoSort(graph map[string]map[string]bool, visited map[string]bool, node string, sorted *[]string) {
+	visited[node] = true
+	for child := range graph[node] {
+		if visited[child] {
+			continue
+		}
+		topoSort(graph, visited, child, sorted)
+	}
+	*sorted = append(*sorted, node)
+}
+
+// Orders types in order such that pgtype.Map can successfully register them
+func sortedCompositeTypes(compositeTypes map[string]CompositeType) []CompositeType {
+	// Map of composite type names to a set of every child composite type name
+	graph := make(map[string]map[string]bool)
+	for typeName, ct := range compositeTypes {
+		graph[typeName] = make(map[string]bool)
+		for _, field := range ct.Fields {
+			fieldType := trimSliceAndPointerPrefix(field.Type)
+			if _, ok := compositeTypes[fieldType]; ok {
+				graph[typeName][fieldType] = true
+			}
+		}
+	}
+
+	visited := make(map[string]bool)
+	sorted := []string{}
+	for typeName := range compositeTypes {
+		if visited[typeName] {
+			continue
+		}
+		topoSort(graph, visited, typeName, &sorted)
+	}
+
+	compositeTypeArr := []CompositeType{}
+	for _, typeName := range sorted {
+		compositeTypeArr = append(compositeTypeArr, compositeTypes[typeName])
+	}
+	return compositeTypeArr
+}
+
+func buildCompositeTypes(req *plugin.GenerateRequest, options *opts.Options) []CompositeType {
+	compositeTypes := make(map[string]CompositeType)
+	for _, schema := range req.Catalog.Schemas {
+		if schema.Name == "pg_catalog" || schema.Name == "information_schema" {
+			continue
+		}
+		for _, ct := range schema.CompositeTypes {
+			var typeName string
+			var sqlName string
+			if schema.Name == req.Catalog.DefaultSchema {
+				typeName = ct.Name
+				sqlName = ct.Name
+			} else {
+				typeName = schema.Name + "_" + ct.Name
+				sqlName = schema.Name + "." + ct.Name
+			}
+			typeName = StructName(typeName, options)
+			compositeType := CompositeType{
+				SQLName: sqlName,
+				Name:    typeName,
+				Comment: ct.Comment,
+			}
+			for _, column := range ct.Columns {
+				compositeType.Fields = append(compositeType.Fields, Field{
+					Name:    StructName(column.Name, options),
+					Type:    goType(req, options, column),
+					Comment: column.Comment,
+				})
+			}
+			compositeTypes[typeName] = compositeType
+		}
+	}
+	return sortedCompositeTypes(compositeTypes)
+}
+
 func buildStructs(req *plugin.GenerateRequest, options *opts.Options) []Struct {
 	var structs []Struct
 	for _, schema := range req.Catalog.Schemas {