snapshot of https://github.com/sqlc-dev/sqlc.git (main, 2026-09-16; blobs > 131072B omitted)

internal/codegen/golang/gen.go

@@ -0,0 +1,426 @@
+package golang
+
+import (
+	"bufio"
+	"bytes"
+	"context"
+	"errors"
+	"fmt"
+	"go/format"
+	"slices"
+	"strings"
+	"text/template"
+
+	"github.com/sqlc-dev/sqlc/internal/codegen/golang/opts"
+	"github.com/sqlc-dev/sqlc/internal/codegen/sdk"
+	"github.com/sqlc-dev/sqlc/internal/metadata"
+	"github.com/sqlc-dev/sqlc/internal/plugin"
+)
+
+type tmplCtx struct {
+	Q             string
+	Package       string
+	ModelsPackage string
+	SQLDriver     opts.SQLDriver
+	Enums         []Enum
+	Structs       []Struct
+	GoQueries     []Query
+	SqlcVersion   string
+
+	// TODO: Race conditions
+	SourceName string
+
+	EmitJSONTags              bool
+	JsonTagsIDUppercase       bool
+	EmitDBTags                bool
+	EmitPreparedQueries       bool
+	EmitInterface             bool
+	EmitEmptySlices           bool
+	EmitMethodsWithDBArgument bool
+	EmitEnumValidMethod       bool
+	EmitAllEnumValues         bool
+	UsesCopyFrom              bool
+	UsesBatch                 bool
+	OmitSqlcVersion           bool
+	BuildTags                 string
+	WrapErrors                bool
+}
+
+func (t *tmplCtx) OutputQuery(sourceName string) bool {
+	return t.SourceName == sourceName
+}
+
+func (t *tmplCtx) codegenDbarg() string {
+	if t.EmitMethodsWithDBArgument {
+		return "db DBTX, "
+	}
+	return ""
+}
+
+// Called as a global method since subtemplate queryCodeStdExec does not have
+// access to the toplevel tmplCtx
+func (t *tmplCtx) codegenEmitPreparedQueries() bool {
+	return t.EmitPreparedQueries
+}
+
+func (t *tmplCtx) codegenQueryMethod(q Query) string {
+	db := "q.db"
+	if t.EmitMethodsWithDBArgument {
+		db = "db"
+	}
+
+	switch q.Cmd {
+	case ":one":
+		if t.EmitPreparedQueries {
+			return "q.queryRow"
+		}
+		return db + ".QueryRowContext"
+
+	case ":many":
+		if t.EmitPreparedQueries {
+			return "q.query"
+		}
+		return db + ".QueryContext"
+
+	default:
+		if t.EmitPreparedQueries {
+			return "q.exec"
+		}
+		return db + ".ExecContext"
+	}
+}
+
+func (t *tmplCtx) codegenQueryRetval(q Query) (string, error) {
+	switch q.Cmd {
+	case ":one":
+		return "row :=", nil
+	case ":many":
+		return "rows, err :=", nil
+	case ":exec":
+		return "_, err :=", nil
+	case ":execrows", ":execlastid":
+		return "result, err :=", nil
+	case ":execresult":
+		if t.WrapErrors {
+			return "result, err :=", nil
+		}
+		return "return", nil
+	default:
+		return "", fmt.Errorf("unhandled q.Cmd case %q", q.Cmd)
+	}
+}
+
+func Generate(ctx context.Context, req *plugin.GenerateRequest) (*plugin.GenerateResponse, error) {
+	options, err := opts.Parse(req)
+	if err != nil {
+		return nil, err
+	}
+
+	if err := opts.ValidateOpts(options); err != nil {
+		return nil, err
+	}
+
+	enums := buildEnums(req, options)
+	structs := buildStructs(req, options)
+	queries, err := buildQueries(req, options, enums, structs)
+	if err != nil {
+		return nil, err
+	}
+
+	if options.OmitUnusedStructs {
+		enums, structs = filterUnusedStructs(enums, structs, queries, options.ModelsTypeQualifier())
+	}
+
+	if err := validate(options, enums, structs, queries); err != nil {
+		return nil, err
+	}
+
+	return generate(req, options, enums, structs, queries)
+}
+
+func validate(options *opts.Options, enums []Enum, structs []Struct, queries []Query) error {
+	enumNames := make(map[string]struct{})
+	for _, enum := range enums {
+		enumNames[enum.Name] = struct{}{}
+		enumNames["Null"+enum.Name] = struct{}{}
+	}
+	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)
+		}
+		structNames[struckt.Name] = struct{}{}
+	}
+	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)
+		}
+	}
+	return nil
+}
+
+func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum, structs []Struct, queries []Query) (*plugin.GenerateResponse, error) {
+	i := &importer{
+		Options: options,
+		Queries: queries,
+		Enums:   enums,
+		Structs: structs,
+	}
+
+	tctx := tmplCtx{
+		EmitInterface:             options.EmitInterface,
+		EmitJSONTags:              options.EmitJsonTags,
+		JsonTagsIDUppercase:       options.JsonTagsIdUppercase,
+		EmitDBTags:                options.EmitDbTags,
+		EmitPreparedQueries:       options.EmitPreparedQueries,
+		EmitEmptySlices:           options.EmitEmptySlices,
+		EmitMethodsWithDBArgument: options.EmitMethodsWithDbArgument,
+		EmitEnumValidMethod:       options.EmitEnumValidMethod,
+		EmitAllEnumValues:         options.EmitAllEnumValues,
+		UsesCopyFrom:              usesCopyFrom(queries),
+		UsesBatch:                 usesBatch(queries),
+		SQLDriver:                 parseDriver(options.SqlPackage),
+		Q:                         "`",
+		Package:                   options.Package,
+		ModelsPackage:             options.ModelsPackage(),
+		Enums:                     enums,
+		Structs:                   structs,
+		SqlcVersion:               req.SqlcVersion,
+		BuildTags:                 options.BuildTags,
+		OmitSqlcVersion:           options.OmitSqlcVersion,
+		WrapErrors:                options.WrapErrors,
+	}
+
+	if tctx.UsesCopyFrom && !tctx.SQLDriver.IsPGX() && options.SqlDriver != opts.SQLDriverGoSQLDriverMySQL {
+		return nil, errors.New(":copyfrom is only supported by pgx and github.com/go-sql-driver/mysql")
+	}
+
+	if tctx.UsesCopyFrom && options.SqlDriver == opts.SQLDriverGoSQLDriverMySQL {
+		if err := checkNoTimesForMySQLCopyFrom(queries); err != nil {
+			return nil, err
+		}
+		tctx.SQLDriver = opts.SQLDriverGoSQLDriverMySQL
+	}
+
+	if tctx.UsesBatch && !tctx.SQLDriver.IsPGX() {
+		return nil, errors.New(":batch* commands are only supported by pgx")
+	}
+
+	funcMap := template.FuncMap{
+		"lowerTitle": sdk.LowerTitle,
+		"comment":    sdk.DoubleSlashComment,
+		"escape":     sdk.EscapeBacktick,
+		"imports":    i.Imports,
+		"hasImports": i.HasImports,
+		"hasPrefix":  strings.HasPrefix,
+
+		// These methods are Go specific, they do not belong in the codegen package
+		// (as that is language independent)
+		"dbarg":               tctx.codegenDbarg,
+		"emitPreparedQueries": tctx.codegenEmitPreparedQueries,
+		"queryMethod":         tctx.codegenQueryMethod,
+		"queryRetval":         tctx.codegenQueryRetval,
+	}
+
+	tmpl := template.Must(
+		template.New("table").
+			Funcs(funcMap).
+			ParseFS(
+				templates,
+				"templates/*.tmpl",
+				"templates/*/*.tmpl",
+			),
+	)
+
+	output := map[string]string{}
+
+	execute := func(name, templateName string) error {
+		imports := i.Imports(name)
+		replacedQueries := replaceConflictedArg(imports, queries)
+
+		var b bytes.Buffer
+		w := bufio.NewWriter(&b)
+		tctx.SourceName = name
+		tctx.GoQueries = replacedQueries
+		err := tmpl.ExecuteTemplate(w, templateName, &tctx)
+		w.Flush()
+		if err != nil {
+			return err
+		}
+		code, err := format.Source(b.Bytes())
+		if err != nil {
+			fmt.Println(b.String())
+			return fmt.Errorf("source error: %w", err)
+		}
+
+		if templateName == "queryFile" && options.OutputFilesSuffix != "" {
+			name += options.OutputFilesSuffix
+		}
+
+		if !strings.HasSuffix(name, ".go") {
+			name += ".go"
+		}
+		output[name] = string(code)
+		return nil
+	}
+
+	dbFileName := "db.go"
+	if options.OutputDbFileName != "" {
+		dbFileName = options.OutputDbFileName
+	}
+	modelsFileName := "models.go"
+	if options.OutputModelsFileName != "" {
+		modelsFileName = options.OutputModelsFileName
+	}
+	querierFileName := "querier.go"
+	if options.OutputQuerierFileName != "" {
+		querierFileName = options.OutputQuerierFileName
+	}
+	copyfromFileName := "copyfrom.go"
+	if options.OutputCopyfromFileName != "" {
+		copyfromFileName = options.OutputCopyfromFileName
+	}
+
+	batchFileName := "batch.go"
+	if options.OutputBatchFileName != "" {
+		batchFileName = options.OutputBatchFileName
+	}
+
+	if err := execute(dbFileName, "dbFile"); err != nil {
+		return nil, err
+	}
+	if options.ModelsEmitEnabled() {
+		if err := execute(modelsFileName, "modelsFile"); err != nil {
+			return nil, err
+		}
+	}
+	if options.EmitInterface {
+		if err := execute(querierFileName, "interfaceFile"); err != nil {
+			return nil, err
+		}
+	}
+	if tctx.UsesCopyFrom {
+		if err := execute(copyfromFileName, "copyfromFile"); err != nil {
+			return nil, err
+		}
+	}
+	if tctx.UsesBatch {
+		if err := execute(batchFileName, "batchFile"); err != nil {
+			return nil, err
+		}
+	}
+
+	files := map[string]struct{}{}
+	for _, gq := range queries {
+		files[gq.SourceName] = struct{}{}
+	}
+
+	for source := range files {
+		if err := execute(source, "queryFile"); err != nil {
+			return nil, err
+		}
+	}
+	resp := plugin.GenerateResponse{}
+
+	for filename, code := range output {
+		resp.Files = append(resp.Files, &plugin.File{
+			Name:     filename,
+			Contents: []byte(code),
+		})
+	}
+
+	return &resp, nil
+}
+
+func usesCopyFrom(queries []Query) bool {
+	for _, q := range queries {
+		if q.Cmd == metadata.CmdCopyFrom {
+			return true
+		}
+	}
+	return false
+}
+
+func usesBatch(queries []Query) bool {
+	for _, q := range queries {
+		if slices.Contains([]string{metadata.CmdBatchExec, metadata.CmdBatchMany, metadata.CmdBatchOne}, q.Cmd) {
+			return true
+		}
+	}
+	return false
+}
+
+func checkNoTimesForMySQLCopyFrom(queries []Query) error {
+	for _, q := range queries {
+		if q.Cmd != metadata.CmdCopyFrom {
+			continue
+		}
+		for _, f := range q.Arg.CopyFromMySQLFields() {
+			if f.Type == "time.Time" {
+				return fmt.Errorf("values with a timezone are not yet supported")
+			}
+		}
+	}
+	return nil
+}
+
+func filterUnusedStructs(enums []Enum, structs []Struct, queries []Query, qualifier string) ([]Enum, []Struct) {
+	keepTypes := make(map[string]struct{})
+
+	keep := func(t string) {
+		keepTypes[t] = struct{}{}
+		// Also store the bare type name so that lookups against
+		// bare struct/enum names match even when types have been
+		// qualified with the models package prefix (e.g. "model.User").
+		if bare := stripQualifier(t, qualifier); bare != t {
+			keepTypes[bare] = struct{}{}
+		}
+	}
+
+	for _, query := range queries {
+		if !query.Arg.isEmpty() {
+			keep(query.Arg.Type())
+			if query.Arg.IsStruct() {
+				for _, field := range query.Arg.Struct.Fields {
+					keep(field.Type)
+				}
+			}
+		}
+		if query.hasRetType() {
+			keep(query.Ret.Type())
+			if query.Ret.IsStruct() {
+				for _, field := range query.Ret.Struct.Fields {
+					keep(strings.TrimPrefix(field.Type, "[]"))
+					for _, embedField := range field.EmbedFields {
+						keep(embedField.Type)
+					}
+				}
+			}
+		}
+	}
+
+	keepEnums := make([]Enum, 0, len(enums))
+	for _, enum := range enums {
+		_, keep := keepTypes[enum.Name]
+		_, keepNull := keepTypes["Null"+enum.Name]
+		_, keepPointer := keepTypes["*"+enum.Name]
+		if keep || keepNull || keepPointer {
+			keepEnums = append(keepEnums, enum)
+		}
+	}
+
+	keepStructs := make([]Struct, 0, len(structs))
+	for _, st := range structs {
+		if _, ok := keepTypes[st.Name]; ok {
+			keepStructs = append(keepStructs, st)
+		}
+	}
+
+	return keepEnums, keepStructs
+}