Untested golang codegen for pgx v5

eaglesemanation · eaglesemanation · commit 69ccc9472cd8 · 2025-04-04T01:25:24.000-05:00
diff --git a/internal/codegen/golang/composite_type.go b/internal/codegen/golang/composite_type.go
@@ -0,0 +1,8 @@
+package golang
+
+type CompositeType struct {
+	SQLName string
+	Name    string
+	Comment string
+	Fields  []Field
+}
diff --git a/internal/codegen/golang/gen.go b/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,
diff --git a/internal/codegen/golang/imports.go b/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,9 @@ 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"})
+		}
 	default:
 		std = append(std, ImportSpec{Path: "database/sql"})
 		if i.Options.EmitPreparedQueries {
diff --git a/internal/codegen/golang/opts/enum.go b/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
 }
diff --git a/internal/codegen/golang/postgresql_type.go b/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
 				}
 			}
 		}
diff --git a/internal/codegen/golang/result.go b/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 {
diff --git a/internal/codegen/golang/templates/pgx/dbCode.tmpl b/internal/codegen/golang/templates/pgx/dbCode.tmpl
@@ -10,6 +10,10 @@ type DBTX interface {
 {{- if .UsesBatch }}
 	SendBatch(context.Context, *pgx.Batch) pgx.BatchResults
 {{- end }}
+{{- if gt (len .CompositeTypes) 0 }}
+	LoadTypes(ctx context.Context, typeNames []string) ([]*pgtype.Type, error)
+	TypeMap() *pgtype.Map
+{{- end }}
 }
 
 {{ if .EmitMethodsWithDBArgument}}
@@ -27,6 +31,26 @@ type Queries struct {
     {{end}}
 }
 
+{{if gt (len .CompositeTypes) 0}}
+{{if not .EmitMethodsWithDBArgument}}
+func (q *Queries) RegisterTypes(ctx context.Context) error {
+	db := q.db
+{{else}}
+func (q *Queries) RegisterTypes(ctx context.Context, db *Queries) error {
+{{end}}
+	typeNames := []string{ 
+	{{- range .CompositeTypes}}
+		"{{.SQLName}}",
+	{{- end}}
+	}
+	dataTypes, err := db.LoadTypes(ctx, typeNames)
+	if err != nil {
+		return err
+	}
+	db.TypeMap().RegisterTypes(dataTypes)
+	return nil
+}
+{{end}}
 {{if not .EmitMethodsWithDBArgument}}
 func (q *Queries) WithTx(tx pgx.Tx) *Queries {
 	return &Queries{
diff --git a/internal/codegen/golang/templates/template.tmpl b/internal/codegen/golang/templates/template.tmpl
@@ -150,6 +150,17 @@ func All{{ .Name }}Values() []{{ .Name }} {
 {{ end }}
 {{end}}
 
+{{range .CompositeTypes}}
+{{if .Comment}}{{comment .Comment}}{{end}}
+type {{.Name}} struct { {{- range .Fields}}
+  {{- if .Comment}}
+  {{comment .Comment}}{{else}}
+  {{- end}}
+  {{.Name}} {{.Type}} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
+  {{- end}}
+}
+{{end}}
+
 {{range .Structs}}
 {{if .Comment}}{{comment .Comment}}{{end}}
 type {{.Name}} struct { {{- range .Fields}}

Original file line number	Diff line number	Diff line change
`@@ -48,6 +48,10 @@ func (d SQLDriver) IsPGX() bool {`
`48`	`48`	`return d == SQLDriverPGXV4 \|\| d == SQLDriverPGXV5`
`49`	`49`	`}`
`50`	`50`
	`51`	`+func (d SQLDriver) IsPGXV5() bool {`
	`52`	`+ return d == SQLDriverPGXV5`
	`53`	`+}`
	`54`	`+`
`51`	`55`	`func (d SQLDriver) IsGoSQLDriverMySQL() bool {`
`52`	`56`	`return d == SQLDriverGoSQLDriverMySQL`
`53`	`57`	`}`