feat: add sqlc.embed to allow model re-use

Author: nickjackson

internal/cmd/shim.go

@@ -241,6 +241,8 @@ func pluginQueryColumn(c *compiler.Column) *plugin.Column {
 		Length:       int32(l),
 		IsNamedParam: c.IsNamedParam,
 		IsFuncCall:   c.IsFuncCall,
+		Embedded:     c.Embedded,
+		EmbedIndex:   c.EmbedIndex,
 	}
 
 	if c.Type != nil {

internal/codegen/golang/query.go

@@ -134,6 +134,17 @@ func (v QueryValue) Scan() string {
 			out = append(out, "&"+v.Name)
 		}
 	} else {
+
+		for _, e := range v.Struct.EmbedFields {
+			for _, f := range e.SubFields {
+				if strings.HasPrefix(f.Type, "[]") && f.Type != "[]byte" && v.SQLPackage != SQLPackagePGX {
+					out = append(out, "pq.Array(&"+v.Name+"."+e.Name+"."+f.Name+")")
+				} else {
+					out = append(out, "&"+v.Name+"."+e.Name+"."+f.Name)
+				}
+			}
+		}
+
 		for _, f := range v.Struct.Fields {
 			if strings.HasPrefix(f.Type, "[]") && f.Type != "[]byte" && !v.SQLDriver.IsPGX() {
 				out = append(out, "pq.Array(&"+v.Name+"."+f.Name+")")
@@ -146,6 +157,7 @@ func (v QueryValue) Scan() string {
 		return strings.Join(out, ",")
 	}
 	out = append(out, "")
+
 	return "\n" + strings.Join(out, ",\n")
 }
 

internal/codegen/golang/result.go

@@ -105,6 +105,11 @@ type goColumn struct {
 	*plugin.Column
 }
 
+type goEmbed struct {
+	id    int
+	Model *Struct
+}
+
 func columnName(c *plugin.Column, pos int) string {
 	if c.Name != "" {
 		return c.Name
@@ -177,7 +182,7 @@ func buildQueries(req *plugin.CodeGenRequest, structs []Struct) ([]Query, error)
 					Column: p.Column,
 				})
 			}
-			s, err := columnsToStruct(req, gq.MethodName+"Params", cols, false)
+			s, err := columnsToStruct(req, gq.MethodName+"Params", nil, cols, false)
 			if err != nil {
 				return nil, err
 			}
@@ -202,48 +207,16 @@ func buildQueries(req *plugin.CodeGenRequest, structs []Struct) ([]Query, error)
 				SQLDriver: sqlpkg,
 			}
 		} else if putOutColumns(query) {
-			var gs *Struct
-			var emit bool
 
-			for _, s := range structs {
-				if len(s.Fields) != len(query.Columns) {
-					continue
-				}
-				same := true
-				for i, f := range s.Fields {
-					c := query.Columns[i]
-					sameName := f.Name == StructName(columnName(c, i), req.Settings)
-					sameType := f.Type == goType(req, c)
-					sameTable := sdk.SameTableName(c.Table, &s.Table, req.Catalog.DefaultSchema)
-					if !sameName || !sameType || !sameTable {
-						same = false
-					}
-				}
-				if same {
-					gs = &s
-					break
-				}
+			qvs, emit, err := buildReturnStruct(req, structs, gq, query.Columns)
+			if err != nil {
+				return nil, err
 			}
 
-			if gs == nil {
-				var columns []goColumn
-				for i, c := range query.Columns {
-					columns = append(columns, goColumn{
-						id:     i,
-						Column: c,
-					})
-				}
-				var err error
-				gs, err = columnsToStruct(req, gq.MethodName+"Row", columns, true)
-				if err != nil {
-					return nil, err
-				}
-				emit = true
-			}
 			gq.Ret = QueryValue{
 				Emit:        emit,
 				Name:        "i",
-				Struct:      gs,
+				Struct:      qvs,
 				SQLDriver:   sqlpkg,
 				EmitPointer: req.Settings.Go.EmitResultStructPointers,
 			}
@@ -267,6 +240,46 @@ func putOutColumns(query *plugin.Query) bool {
 	return false
 }
 
+func buildReturnStruct(req *plugin.CodeGenRequest, models []Struct, gq Query, columns []*plugin.Column) (*Struct, bool, error) {
+	// group columns by embeds or top level fields
+	rootCols, embedGroups := groupColumns(columns)
+
+	// no embeds, just top level fields
+	if len(embedGroups) == 0 {
+		// return early if model already exists
+		if model := lookupModelForColumns(req, models, rootCols); model != nil {
+			return model, false, nil
+		}
+	}
+
+	// reaching here means we need a to construct a new result struct, which
+	// may or may not contain embeds
+	var goCols []goColumn
+	for i, c := range rootCols {
+		goCols = append(goCols, goColumn{
+			id:     i,
+			Column: c,
+		})
+	}
+
+	var goEmbeds []goEmbed
+	for i, eg := range embedGroups {
+		if model := lookupModelForColumns(req, models, eg.columns); model != nil {
+			goEmbeds = append(goEmbeds, goEmbed{
+				id:    i,
+				Model: model,
+			})
+		}
+	}
+
+	gs, err := columnsToStruct(req, gq.MethodName+"Row", goEmbeds, goCols, true)
+	if err != nil {
+		return nil, false, err
+	}
+
+	return gs, true, nil
+}
+
 // It's possible that this method will generate duplicate JSON tag values
 //
 //	Columns: count, count,   count_2
@@ -275,12 +288,13 @@ func putOutColumns(query *plugin.Query) bool {
 // JSON tags: count, count_2, count_2
 //
 // This is unlikely to happen, so don't fix it yet
-func columnsToStruct(req *plugin.CodeGenRequest, name string, columns []goColumn, useID bool) (*Struct, error) {
+func columnsToStruct(req *plugin.CodeGenRequest, name string, embeds []goEmbed, columns []goColumn, useID bool) (*Struct, error) {
 	gs := Struct{
 		Name: name,
 	}
 	seen := map[string][]int{}
 	suffixes := map[int]int{}
+
 	for i, c := range columns {
 		colName := columnName(c.Column, i)
 		tagName := colName
@@ -342,6 +356,42 @@ func columnsToStruct(req *plugin.CodeGenRequest, name string, columns []goColumn
 		return nil, err
 	}
 
+	seenEmbeds := map[string]int{}
+
+	for _, em := range embeds {
+
+		fieldName := em.Model.Name
+		tagName := fieldName
+
+		seen, ok := seenEmbeds[fieldName]
+		if !ok {
+			seenEmbeds[fieldName] = 0
+		}
+		seen++
+		seenEmbeds[fieldName] = seen
+
+		suffix := seen - 1
+		if suffix > 0 {
+			tagName = fmt.Sprintf("%s_%d", tagName, suffix)
+			fieldName = fmt.Sprintf("%s_%d", fieldName, suffix)
+		}
+
+		tags := map[string]string{}
+
+		if req.Settings.Go.EmitJsonTags {
+			tags["json:"] = JSONTagName(tagName, req.Settings)
+		}
+
+		gs.EmbedFields = append(gs.EmbedFields, EmbedField{
+			Field: Field{
+				Name: fieldName,
+				Type: em.Model.Name,
+				Tags: tags,
+			},
+			SubFields: em.Model.Fields,
+		})
+	}
+
 	return &gs, nil
 }
 
@@ -356,3 +406,61 @@ func checkIncompatibleFieldTypes(fields []Field) error {
 	}
 	return nil
 }
+
+type embedGroup struct {
+	index   int32
+	columns []*plugin.Column
+}
+
+func groupColumns(cols []*plugin.Column) ([]*plugin.Column, []*embedGroup) {
+	root := []*plugin.Column{}
+	embeds := []*embedGroup{}
+
+	for _, c := range cols {
+
+		if !c.Embedded {
+			root = append(root, c)
+			continue
+		}
+
+		var found bool
+		for _, v := range embeds {
+			if v.index == c.EmbedIndex {
+				found = true
+				v.columns = append(v.columns, c)
+			}
+		}
+		if !found {
+			embeds = append(embeds, &embedGroup{
+				index:   c.EmbedIndex,
+				columns: []*plugin.Column{c},
+			})
+		}
+	}
+
+	return root, embeds
+}
+
+func lookupModelForColumns(req *plugin.CodeGenRequest, models []Struct, columns []*plugin.Column) *Struct {
+	for _, model := range models {
+		if len(model.Fields) != len(columns) {
+			continue
+		}
+
+		same := true
+		for i, f := range model.Fields {
+			c := columns[i]
+			sameName := f.Name == StructName(columnName(c, i), req.Settings)
+			sameType := f.Type == goType(req, c)
+			sameTable := sdk.SameTableName(c.Table, &model.Table, req.Catalog.DefaultSchema)
+			if !sameName || !sameType || !sameTable {
+				same = false
+			}
+		}
+		if same {
+			return &model
+		}
+	}
+
+	return nil
+}

internal/codegen/golang/struct.go

@@ -9,10 +9,16 @@ import (
 )
 
 type Struct struct {
-	Table   plugin.Identifier
-	Name    string
-	Fields  []Field
-	Comment string
+	Table       plugin.Identifier
+	Name        string
+	EmbedFields []EmbedField
+	Fields      []Field
+	Comment     string
+}
+
+type EmbedField struct {
+	Field
+	SubFields []Field
 }
 
 func StructName(name string, settings *plugin.Settings) string {

internal/codegen/golang/templates/pgx/queryCode.tmpl

@@ -16,7 +16,10 @@ type {{.Arg.Type}} struct { {{- range .Arg.Struct.Fields}}
 {{end}}
 
 {{if .Ret.EmitStruct}}
-type {{.Ret.Type}} struct { {{- range .Ret.Struct.Fields}}
+type {{.Ret.Type}} struct { {{- range .Ret.Struct.EmbedFields}}
+  {{.Name}} {{.Type}} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
+  {{- end}}
+  {{- range .Ret.Struct.Fields}}
   {{.Name}} {{.Type}} {{if .Tag}}{{$.Q}}{{.Tag}}{{$.Q}}{{end}}
   {{- end}}
 }

internal/compiler/expand.go

@@ -132,9 +132,16 @@ func (c *Compiler) expandStmt(qc *QueryCatalog, raw *ast.RawStmt, node ast.Node)
 		for _, p := range parts {
 			old = append(old, c.quoteIdent(p))
 		}
+		oldString := strings.Join(old, ".")
+
+		// use the sqlc.embed string instead
+		if embed, ok := qc.embeds.Find(ref); ok {
+			oldString = embed.Orig()
+		}
+
 		edits = append(edits, source.Edit{
 			Location: res.Location - raw.StmtLocation,
-			Old:      strings.Join(old, "."),
+			Old:      oldString,
 			New:      strings.Join(cols, ", "),
 		})
 	}

internal/compiler/output_columns.go

@@ -14,7 +14,7 @@ import (
 
 // OutputColumns determines which columns a statement will output
 func (c *Compiler) OutputColumns(stmt ast.Node) ([]*catalog.Column, error) {
-	qc, err := buildQueryCatalog(c.catalog, stmt)
+	qc, err := buildQueryCatalog(c.catalog, stmt, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -178,6 +178,9 @@ func outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, error) {
 
 		case *ast.ColumnRef:
 			if hasStarRef(n) {
+
+				embed, isEmbedded := qc.embeds.Find(n)
+
 				// TODO: This code is copied in func expand()
 				for _, t := range tables {
 					scope := astutils.Join(n.Fields, ".")
@@ -189,7 +192,7 @@ func outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, error) {
 						if res.Name != nil {
 							cname = *res.Name
 						}
-						cols = append(cols, &Column{
+						col := &Column{
 							Name:       cname,
 							Type:       c.Type,
 							Scope:      scope,
@@ -199,7 +202,14 @@ func outputColumns(qc *QueryCatalog, node ast.Node) ([]*Column, error) {
 							NotNull:    c.NotNull,
 							IsArray:    c.IsArray,
 							Length:     c.Length,
-						})
+							Embedded:   isEmbedded,
+						}
+
+						if isEmbedded {
+							col.EmbedIndex = embed.Index
+						}
+
+						cols = append(cols, col)
 					}
 				}
 				continue

internal/compiler/parse.go

@@ -83,7 +83,10 @@ func (c *Compiler) parseQuery(stmt ast.Node, src string, o opts.Parser) (*Query,
 	} else {
 		sort.Slice(refs, func(i, j int) bool { return refs[i].ref.Number < refs[j].ref.Number })
 	}
-	qc, err := buildQueryCatalog(c.catalog, raw.Stmt)
+
+	raw, embeds := rewrite.Embeds(c.conf.Engine, raw)
+
+	qc, err := buildQueryCatalog(c.catalog, raw.Stmt, embeds)
 	if err != nil {
 		return nil, err
 	}

internal/compiler/query.go

@@ -30,6 +30,9 @@ type Column struct {
 	TableAlias string
 	Type       *ast.TypeName
 
+	Embedded   bool
+	EmbedIndex int32
+
 	skipTableRequiredCheck bool
 }