compiler: expand stars in the query text on the core path

internal/compiler/engine.go

@@ -38,6 +38,13 @@ type Compiler struct {
 	coreAnalysis bool
 	coreDialect  core.Option
 
+	// newParser builds a parser for the configured engine, and is set for
+	// every engine either path supports. The core path analyzes statements
+	// concurrently, and a parser holds enough state that two goroutines cannot
+	// share one, so a goroutine that has to parse something of its own — the
+	// query text it just rewrote — builds its own rather than taking c.parser.
+	newParser func() Parser
+
 	schema []string
 }
 
@@ -76,7 +83,7 @@ func NewCompiler(conf config.SQL, combo config.CombinedSettings, parserOpts opts
 
 	switch conf.Engine {
 	case config.EngineSQLite:
-		c.parser = sqlite.NewParser()
+		c.newParser = func() Parser { return sqlite.NewParser() }
 		c.catalog = sqlite.NewCatalog()
 		c.selector = newSQLiteSelector()
 
@@ -90,11 +97,11 @@ func NewCompiler(conf config.SQL, combo config.CombinedSettings, parserOpts opts
 			}
 		}
 	case config.EngineMySQL:
-		c.parser = dolphin.NewParser()
+		c.newParser = func() Parser { return dolphin.NewParser() }
 		c.catalog = dolphin.NewCatalog()
 		c.selector = newDefaultSelector()
 	case config.EnginePostgreSQL:
-		c.parser = postgresql.NewParser()
+		c.newParser = func() Parser { return postgresql.NewParser() }
 		c.catalog = postgresql.NewCatalog()
 		c.selector = newDefaultSelector()
 
@@ -110,6 +117,7 @@ func NewCompiler(conf config.SQL, combo config.CombinedSettings, parserOpts opts
 	default:
 		return nil, fmt.Errorf("unknown engine: %s", conf.Engine)
 	}
+	c.parser = c.newParser()
 	return c, nil
 }
 
@@ -119,28 +127,29 @@ func (c *Compiler) initCore() error {
 	var dialect core.Option
 	switch c.conf.Engine {
 	case config.EngineSQLite:
-		c.parser = sqlite.NewParser()
+		c.newParser = func() Parser { return sqlite.NewParser() }
 		c.selector = newSQLiteSelector()
 		dialect = sqlite.Dialect()
 	case config.EngineMySQL:
-		c.parser = dolphin.NewParser()
+		c.newParser = func() Parser { return dolphin.NewParser() }
 		c.selector = newDefaultSelector()
 		dialect = dolphin.Dialect()
 	case config.EnginePostgreSQL:
-		c.parser = postgresql.NewParser()
+		c.newParser = func() Parser { return postgresql.NewParser() }
 		c.selector = newDefaultSelector()
 		dialect = postgresql.Dialect()
 	case config.EngineClickHouse:
-		c.parser = clickhouse.NewParser()
+		c.newParser = func() Parser { return clickhouse.NewParser() }
 		c.selector = newDefaultSelector()
 		dialect = clickhouse.Dialect()
 	case config.EngineGoogleSQL:
-		c.parser = googlesql.NewParser()
+		c.newParser = func() Parser { return googlesql.NewParser() }
 		c.selector = newDefaultSelector()
 		dialect = googlesql.Dialect()
 	default:
 		return fmt.Errorf("unknown engine: %s", c.conf.Engine)
 	}
+	c.parser = c.newParser()
 	c.coreDialect = dialect
 	return nil
 }

internal/compiler/expand.go

@@ -78,6 +78,38 @@ func (c *Compiler) quote(x string) string {
 	}
 }
 
+// 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 {
@@ -157,37 +189,9 @@ func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node)
 				cols = append(cols, cname)
 			}
 		}
-		var old []string
-		for _, p := range parts {
-			if p == "*" {
-				old = append(old, p)
-			} else {
-				old = append(old, c.quoteIdent(p))
-			}
-		}
-
-		// An embed was rewritten to "table.*" in the query text, so it is
-		// measured the same way as a star reference the user wrote.
-		oldFunc := 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
-		}
-
 		edits = append(edits, source.Edit{
 			Location: res.Location - raw.StmtLocation,
-			OldFunc:  oldFunc,
+			OldFunc:  c.starOldFunc(parts),
 			New:      strings.Join(cols, ", "),
 		})
 	}

internal/compiler/expand_core.go

@@ -0,0 +1,76 @@
+package compiler
+
+import (
+	"strings"
+
+	"github.com/sqlc-dev/sqlc/internal/core"
+	"github.com/sqlc-dev/sqlc/internal/source"
+	"github.com/sqlc-dev/sqlc/internal/sql/ast"
+)
+
+// expandCore rewrites the stars in a query's text with the columns the core
+// analyzer resolved them to. The analyzer has already walked the statement and
+// its subqueries, so there is nothing left to look up here: what remains is
+// deciding how each name is written, which is the engine's business and not
+// the core's.
+func (c *Compiler) expandCore(raw *ast.RawStmt, stars []core.StarExpansion) []source.Edit {
+	if len(stars) == 0 {
+		return nil
+	}
+	edits := make([]source.Edit, 0, len(stars))
+	seen := make(map[int]bool, len(stars))
+	for _, star := range stars {
+		// The analyzer types an expression again when a later clause refers to
+		// the output name it was given, so "SELECT (SELECT * FROM t) AS x ...
+		// GROUP BY x" reports the star in the subquery once per pass. The
+		// passes agree on what the star covers, and editing the same reference
+		// twice would overlap, so only the first is kept.
+		if seen[star.Location] {
+			continue
+		}
+		seen[star.Location] = true
+
+		// Everything before the star qualifies it: "foo.*" is scoped to foo,
+		// while a bare "*" covers every relation in the FROM clause.
+		scope := strings.Join(star.Fields[:len(star.Fields)-1], ".")
+
+		// An unqualified star that covers more than one relation may name the
+		// same column twice, so those are written with their relation.
+		counts := map[string]int{}
+		if scope == "" {
+			for _, col := range star.Columns {
+				counts[col.Name]++
+			}
+		}
+
+		cols := make([]string, 0, len(star.Columns))
+		for _, col := range star.Columns {
+			cname := col.Name
+			if star.Alias != "" {
+				cname = star.Alias
+			}
+			cname = c.quoteIdent(cname)
+			if scope != "" {
+				cname = c.quoteIdent(scope) + "." + cname
+			}
+			if counts[cname] > 1 {
+				cname = c.quoteIdent(col.Relation) + "." + 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).
+			cols = append(cols, c.selector.ColumnExpr(cname, &Column{
+				Name:     col.Name,
+				DataType: col.DataType,
+			}))
+		}
+
+		edits = append(edits, source.Edit{
+			Location: star.Location - raw.StmtLocation,
+			OldFunc:  c.starOldFunc(star.Fields),
+			New:      strings.Join(cols, ", "),
+		})
+	}
+	return edits
+}

internal/compiler/parse_core.go

@@ -2,6 +2,7 @@ package compiler
 
 import (
 	"errors"
+	"fmt"
 	"strings"
 
 	"github.com/sqlc-dev/sqlc/internal/core"
@@ -64,6 +65,17 @@ func (c *Compiler) parseQueryCore(raw *ast.RawStmt, src string, pre *preprocess.
 		for _, p := range res.Parameters {
 			params = append(params, Parameter{Number: p.Number, Column: coreParamColumn(p, namedParams)})
 		}
+		expanded, err = source.Mutate(rawSQL, c.expandCore(raw, res.Stars))
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	// If the query string was edited, make sure the syntax is valid
+	if expanded != rawSQL {
+		if _, err := c.newParser().Parse(strings.NewReader(expanded)); err != nil {
+			return nil, fmt.Errorf("edited query syntax is invalid: %w", err)
+		}
 	}
 
 	trimmed, comments, err := source.StripComments(expanded)

internal/core/analysis.go

@@ -10,9 +10,39 @@ const (
 )
 
 type PrepareResult struct {
-	Command    Command     `json:"command,omitempty"`
-	Columns    []Column    `json:"columns"`
-	Parameters []Parameter `json:"parameters"`
+	Command    Command         `json:"command,omitempty"`
+	Columns    []Column        `json:"columns"`
+	Parameters []Parameter     `json:"parameters"`
+	Stars      []StarExpansion `json:"stars,omitempty"`
+}
+
+// StarExpansion is what a star in a target list stands for. The analyzer
+// resolves the reference against the query's scope and reports the columns it
+// covers; rewriting the query text with them is the caller's to do, since only
+// it knows how the engine quotes an identifier.
+type StarExpansion struct {
+	// Location is where the target the star belongs to starts, measured the
+	// way the AST measures a node: from the beginning of the file the
+	// statement was parsed from.
+	Location int `json:"location"`
+
+	// Fields is the reference as it was written, with the star as its last
+	// element: ["*"] for a bare star and ["foo", "*"] for a qualified one.
+	Fields []string `json:"fields"`
+
+	// Alias is the output name the target was given, if any.
+	Alias string `json:"alias,omitempty"`
+
+	Columns []StarColumn `json:"columns"`
+}
+
+// StarColumn is a single column a star expanded to.
+type StarColumn struct {
+	// Relation is the name the column's relation goes by in the query, which
+	// is its alias when it was given one.
+	Relation string `json:"relation,omitempty"`
+	Name     string `json:"name"`
+	DataType string `json:"data_type,omitempty"`
 }
 
 type ColumnSource struct {

internal/core/analyzer/analyzer.go

@@ -14,6 +14,7 @@ func Prepare(cat *core.Catalog, stmt ast.Node) (core.PrepareResult, error) {
 	a := &analyzer{
 		cat:    cat,
 		params: map[int]core.Parameter{},
+		stars:  &[]core.StarExpansion{},
 	}
 	switch s := stmt.(type) {
 	case *ast.SelectStmt:
@@ -63,6 +64,18 @@ type analyzer struct {
 
 	// resolving guards against an alias that refers to itself.
 	resolving map[string]bool
+
+	// stars are the expansions every star in the statement asked for, shared
+	// with the analyzers of the queries nested in it so one statement reports
+	// all of them.
+	stars *[]core.StarExpansion
+}
+
+func (a *analyzer) recordStar(s core.StarExpansion) {
+	if a.stars == nil {
+		return
+	}
+	*a.stars = append(*a.stars, s)
 }
 
 // subquery analyzes a nested SELECT. It shares the parameter set, so a
@@ -74,6 +87,7 @@ func (a *analyzer) subquery(s *ast.SelectStmt) (*analyzer, error) {
 		params: a.params,
 		outer:  a.scope,
 		ctes:   a.ctes,
+		stars:  a.stars,
 	}
 	if err := sub.analyzeSelect(s); err != nil {
 		return nil, err
@@ -105,11 +119,15 @@ func derivedRel(alias string, cols []core.Column) scopeRel {
 }
 
 func (a *analyzer) result() core.PrepareResult {
-	return core.PrepareResult{
+	res := core.PrepareResult{
 		Command:    a.command,
 		Columns:    a.columns,
 		Parameters: orderedParams(a.params),
 	}
+	if a.stars != nil {
+		res.Stars = *a.stars
+	}
+	return res
 }
 
 func orderedParams(m map[int]core.Parameter) []core.Parameter {