upstream-sqlc/internal/compiler/expand.go

package compiler

import (
	"fmt"
	"regexp"
	"strings"

	"github.com/sqlc-dev/sqlc/internal/config"
	"github.com/sqlc-dev/sqlc/internal/source"
	"github.com/sqlc-dev/sqlc/internal/sql/ast"
	"github.com/sqlc-dev/sqlc/internal/sql/astutils"
)

func (c *Compiler) expand(qc *QueryCatalog, raw *ast.RawStmt) ([]source.Edit, error) {
	// Return early if there are no A_Star nodes to expand
	stars := astutils.Search(raw, func(node ast.Node) bool {
		_, ok := node.(*ast.A_Star)
		return ok
	})
	if len(stars.Items) == 0 {
		return nil, nil
	}
	list := astutils.Search(raw, func(node ast.Node) bool {
		switch node.(type) {
		case *ast.DeleteStmt:
		case *ast.InsertStmt:
		case *ast.SelectStmt:
		case *ast.UpdateStmt:
		default:
			return false
		}
		return true
	})
	if len(list.Items) == 0 {
		return nil, nil
	}
	var edits []source.Edit
	for _, item := range list.Items {
		edit, err := c.expandStmt(qc, raw, item)
		if err != nil {
			return nil, err
		}
		edits = append(edits, edit...)
	}
	return edits, nil
}

var validPostgresIdent = regexp.MustCompile(`^[a-z_][a-z0-9_$]*$`)

func (c *Compiler) quoteIdent(ident string) string {
	if c.parser.IsReservedKeyword(ident) {
		return c.quote(ident)
	}
	// SQL identifiers and key words must begin with a letter (a-z, but also
	// letters with diacritical marks and non-Latin letters) or an underscore
	// (_). Subsequent characters in an identifier or key word can be letters,
	// underscores, digits (0-9), or dollar signs ($).
	//
	// https://www.postgresql.org/docs/current/sql-syntax-lexical.html#SQL-SYNTAX-IDENTIFIERS
	if c.conf.Engine == config.EnginePostgreSQL {
		// camelCase means the column is also camelCase
		if strings.ToLower(ident) != ident {
			return c.quote(ident)
		}
		if !validPostgresIdent.MatchString(strings.ToLower(ident)) {
			return c.quote(ident)
		}
	}
	return ident
}

func (c *Compiler) quote(x string) string {
	switch c.conf.Engine {
	case config.EngineMySQL:
		return "`" + x + "`"
	default:
		return "\"" + x + "\""
	}
}

// starOldFunc measures how much of the query text a star reference occupies,
// so an edit replaces the reference and nothing else. Each part is measured
// both bare and quoted: an embed was rewritten to "table.*" in the query text,
// preserving the way the user quoted the table, so it is measured the same way
// as a star reference the user wrote.
func (c *Compiler) starOldFunc(parts []string) func(string) int {
	old := make([]string, 0, len(parts))
	for _, p := range parts {
		if p == "*" {
			old = append(old, p)
		} else {
			old = append(old, c.quoteIdent(p))
		}
	}
	return func(s string) int {
		length := 0
		for i, o := range old {
			if hasSeparator := i > 0; hasSeparator {
				length++
			}
			if strings.HasPrefix(s[length:], o) {
				length += len(o)
			} else if quoted := c.quote(o); strings.HasPrefix(s[length:], quoted) {
				length += len(quoted)
			} else {
				length += len(o)
			}
		}
		return length
	}
}

func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node) ([]source.Edit, error) {
	tables, err := c.sourceTables(qc, node)
	if err != nil {
		return nil, err
	}

	// Virtual tables for the OLD and NEW aliases available in a RETURNING
	// clause (PostgreSQL 18)
	rtables, err := c.returningTables(qc, node)
	if err != nil {
		return nil, err
	}

	var targets *ast.List
	switch n := node.(type) {
	case *ast.DeleteStmt:
		targets = n.ReturningList
	case *ast.InsertStmt:
		targets = n.ReturningList
	case *ast.SelectStmt:
		targets = n.TargetList
	case *ast.UpdateStmt:
		targets = n.ReturningList
	default:
		return nil, fmt.Errorf("outputColumns: unsupported node type: %T", n)
	}

	var edits []source.Edit
	for _, target := range targets.Items {
		res, ok := target.(*ast.ResTarget)
		if !ok {
			continue
		}
		ref, ok := res.Val.(*ast.ColumnRef)
		if !ok {
			continue
		}
		if !hasStarRef(ref) {
			continue
		}
		var parts, cols []string
		for _, f := range ref.Fields.Items {
			switch field := f.(type) {
			case *ast.String:
				parts = append(parts, field.Str)
			case *ast.A_Star:
				parts = append(parts, "*")
			default:
				return nil, fmt.Errorf("unknown field in ColumnRef: %T", f)
			}
		}
		scope := astutils.Join(ref.Fields, ".")
		counts := map[string]int{}
		if scope == "" {
			for _, t := range tables {
				for _, c := range t.Columns {
					counts[c.Name] += 1
				}
			}
		}
		starTables := tables
		if vt := returningTableForScope(tables, rtables, scope); vt != nil {
			starTables = []*Table{vt}
		}
		for _, t := range starTables {
			if scope != "" && scope != t.Rel.Name {
				continue
			}
			tableName := c.quoteIdent(t.Rel.Name)
			scopeName := c.quoteIdent(scope)
			for _, column := range t.Columns {
				cname := column.Name
				if res.Name != nil {
					cname = *res.Name
				}
				cname = c.quoteIdent(cname)
				if scope != "" {
					cname = scopeName + "." + cname
				}
				if counts[cname] > 1 {
					cname = tableName + "." + cname
				}

				// This is important for SQLite in particular which needs to
				// wrap jsonb column values with `json(colname)` so they're in a
				// publicly usable format (i.e. not jsonb).
				cname = c.selector.ColumnExpr(cname, column)
				cols = append(cols, cname)
			}
		}
		edits = append(edits, source.Edit{
			Location: res.Location - raw.StmtLocation,
			OldFunc:  c.starOldFunc(parts),
			New:      strings.Join(cols, ", "),
		})
	}

	return edits, nil
}