gen.go

package wicked

import (
	"bufio"
	"bytes"
	"context"
	"errors"
	"fmt"
	"go/format"
	"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
	SQLDriver    opts.SQLDriver
	Enums        []Enum
	Structs      []Struct
	GoQueries    []Query
	SqlcVersion  string
	DumpLoader   *DumpLoader
	RawSchemaSQL string

	// The query file currently being rendered by this generation invocation.
	SourceName string

	EmitJSONTags              bool
	JsonTagsIDUppercase       bool
	EmitDBTags                bool
	EmitPreparedQueries       bool
	EmitInterface             bool
	HasEach                   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) {
	if req == nil || req.Settings == nil || req.Catalog == nil {
		return nil, errors.New("wicked: missing settings or catalog")
	}
	options, err := parseOptions(req)
	if err != nil {
		return nil, err
	}

	if err := opts.ValidateOpts(options); err != nil {
		return nil, err
	}

	if req.GetWicked() == nil || req.GetWicked().PrimaryRelation == nil {
		return nil, errors.New("wicked: missing primary schema metadata")
	}
	if strings.TrimSpace(req.GetWicked().PrimarySchemaSql) == "" {
		return nil, errors.New("wicked: missing primary schema source")
	}
	for _, query := range req.Queries {
		switch query.Cmd {
		case metadata.CmdOne, metadata.CmdMany, metadata.CmdExec, metadata.CmdExecRows, metadata.CmdExecResult, metadata.CmdCopyFrom:
		default:
			return nil, fmt.Errorf("wicked: %s: unsupported command %s", query.Name, query.Cmd)
		}
	}
	if req.Settings.Engine != "postgresql" {
		return nil, errors.New("wicked: only PostgreSQL is supported")
	}
	enums := buildEnums(req, options)
	structs := buildStructs(req, options)
	queries, err := buildQueries(req, options, structs)
	if err != nil {
		return nil, err
	}

	if options.OmitUnusedStructs {
		enums, structs = filterUnusedStructs(enums, structs, queries)
	}

	if err := buildQueryInvalidates(queries); err != nil {
		return nil, err
	}
	var missing []error
	for _, q := range queries {
		if q.Option.Timeout == 0 {
			missing = append(missing, fmt.Errorf("%s/%s does not have a timeout option", options.Package, q.MethodName))
		}
	}
	if len(missing) != 0 {
		return nil, errors.Join(missing...)
	}
	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 {
	if hasEach(queries) {
		if err := validateEachNames(options, enums, structs, queries); err != nil {
			return err
		}
	}
	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,
	}

	dumploader, err := buildDumpLoader(structs)
	if err != nil {
		return nil, err
	}
	tctx := tmplCtx{
		DumpLoader:                dumploader,
		RawSchemaSQL:              req.GetWicked().PrimarySchemaSql,
		EmitInterface:             false,
		HasEach:                   hasEach(queries),
		EmitJSONTags:              true,
		JsonTagsIDUppercase:       false,
		EmitDBTags:                false,
		EmitPreparedQueries:       false,
		EmitEmptySlices:           false,
		EmitMethodsWithDBArgument: false,
		EmitEnumValidMethod:       true,
		EmitAllEnumValues:         true,
		UsesCopyFrom:              usesCopyFrom(queries),
		UsesBatch:                 usesBatch(queries),
		SQLDriver:                 parseDriver(options.SqlPackage),
		Q:                         "`",
		Package:                   options.Package,
		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 {
			return fmt.Errorf("formatting generated %s: %w", name, 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 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 {
		for _, cmd := range []string{metadata.CmdBatchExec, metadata.CmdBatchMany, metadata.CmdBatchOne} {
			if q.Cmd == 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) ([]Enum, []Struct) {
	keepTypes := make(map[string]struct{})

	for _, query := range queries {
		if !query.Arg.isEmpty() {
			keepTypes[query.Arg.Type()] = struct{}{}
			if query.Arg.IsStruct() {
				for _, field := range query.Arg.Struct.Fields {
					keepTypes[field.Type] = struct{}{}
				}
			}
		}
		if query.hasRetType() {
			keepTypes[query.Ret.Type()] = struct{}{}
			if query.Ret.IsStruct() {
				for _, field := range query.Ret.Struct.Fields {
					keepTypes[strings.TrimPrefix(field.Type, "[]")] = struct{}{}
					for _, embedField := range field.EmbedFields {
						keepTypes[embedField.Type] = struct{}{}
					}
				}
			}
		}
	}

	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
}

func hasEach(queries []Query) bool {
	for _, query := range queries {
		if query.Option.Each {
			return true
		}
	}
	return false
}

// Check only declarations introduced by Each. Existing queries without the
// annotation retain the original naming rules and output.
func validateEachNames(options *opts.Options, enums []Enum, structs []Struct, queries []Query) error {
	methods := map[string]bool{
		"WithTx": true, "WithCache": true, "AsReadOnly": true,
		"UseReplica": true, "GetCache": true, "GetConn": true,
		"Dump": true, "Load": true,
	}
	declarations := make(map[string]string)
	for _, enum := range enums {
		declarations[enum.Name] = "enum"
		declarations["Null"+enum.Name] = "enum"
		for _, constant := range enum.Constants {
			declarations[constant.Name] = "enum constant"
		}
	}
	for _, strct := range structs {
		declarations[strct.Name] = "model"
	}
	for _, query := range queries {
		methods[query.MethodName] = true
		if query.Cmd != metadata.CmdCopyFrom {
			declarations[query.ConstantName] = "query constant"
		}
		if query.Cmd == metadata.CmdOne || query.Cmd == metadata.CmdMany {
			declarations["_"+query.MethodName] = "helper"
		}
		if query.Arg.EmitStruct() {
			declarations[query.Arg.Type()] = "parameter type"
		}
		if query.Ret.EmitStruct() {
			declarations[query.Ret.Type()] = "row type"
		}
	}
	for _, query := range queries {
		if !query.Option.Each {
			continue
		}
		name := query.MethodName + "Each"
		if methods[name] {
			return fmt.Errorf("%s/%s: Each companion name conflicts: %s", options.Package, query.MethodName, name)
		}
		if kind := declarations["_"+name]; kind != "" {
			return fmt.Errorf("%s/%s: Each companion name conflicts with %s: _%s", options.Package, query.MethodName, kind, name)
		}
	}
	if options.EmitInterface {
		if kind := declarations["EachQuerier"]; kind != "" {
			return fmt.Errorf("Each companion name conflicts with %s: EachQuerier", kind)
		}
	}
	return nil
}