feat: support sqlc.embed

Author: kevinvalk

internal/gen.go

@@ -53,6 +53,8 @@ type Field struct{
 Namestring
 TypepyType
 Commentstring
+// EmbedFields contains the embedded fields that require scanning.
+EmbedFields []Field
 }
 
 typeStructstruct{
@@ -105,14 +107,42 @@ func (v QueryValue) RowNode(rowVar string) *pyast.Node{
 call:=&pyast.Call{
 Func: v.Annotation(),
  }
-fori, f:=rangev.Struct.Fields{
-call.Keywords=append(call.Keywords, &pyast.Keyword{
-Arg: f.Name,
-Value: subscriptNode(
+rowIndex:=0// We need to keep track of the index in the row variable.
+for_, f:=rangev.Struct.Fields{
+
+varvalueNode*pyast.Node
+// Check if we are using sqlc.embed, if so we need to create a new object.
+iflen(f.EmbedFields) >0{
+// We keep this separate so we can easily add all arguments.
+embed_call:=&pyast.Call{Func: f.Type.Annotation()}
+
+// Now add all field Initializers for the embedded model that index into the original row.
+fori, embedField:=rangef.EmbedFields{
+embed_call.Keywords=append(embed_call.Keywords, &pyast.Keyword{
+Arg: embedField.Name,
+Value: subscriptNode(
+rowVar,
+constantInt(rowIndex+i),
+ ),
+ })
+ }
+
+valueNode=&pyast.Node{
+Node: &pyast.Node_Call{
+Call: embed_call,
+ },
+ }
+
+rowIndex+=len(f.EmbedFields)
+ } else{
+valueNode=subscriptNode(
 rowVar,
-constantInt(i),
- ),
- })
+constantInt(rowIndex),
+ )
+rowIndex++
+ }
+
+call.Keywords=append(call.Keywords, &pyast.Keyword{Arg: f.Name, Value: valueNode})
  }
 return&pyast.Node{
 Node: &pyast.Node_Call{
@@ -336,6 +366,47 @@ func paramName(p *plugin.Parameter) string{
 typepyColumnstruct{
 idint32
 *plugin.Column
+embed*pyEmbed
+}
+
+typepyEmbedstruct{
+modelTypestring
+modelNamestring
+fields []Field
+}
+
+// Taken from https://github.com/sqlc-dev/sqlc/blob/8c59fbb9938a0bad3d9971fc2c10ea1f83cc1d0b/internal/codegen/golang/result.go#L123-L126
+// look through all the structs and attempt to find a matching one to embed
+// We need the name of the struct and its field names.
+funcnewGoEmbed(embed*plugin.Identifier, structs []Struct, defaultSchemastring) *pyEmbed{
+ifembed==nil{
+returnnil
+ }
+
+for_, s:=rangestructs{
+embedSchema:=defaultSchema
+ifembed.Schema!=""{
+embedSchema=embed.Schema
+ }
+
+// compare the other attributes
+ifembed.Catalog!=s.Table.Catalog||embed.Name!=s.Table.Name||embedSchema!=s.Table.Schema{
+continue
+ }
+
+fields:=make([]Field, len(s.Fields))
+fori, f:=ranges.Fields{
+fields[i] =f
+ }
+
+return&pyEmbed{
+modelType: s.Name,
+modelName: s.Name,
+fields: fields,
+ }
+ }
+
+returnnil
 }
 
 funccolumnsToStruct(req*plugin.CodeGenRequest, namestring, columns []pyColumn) *Struct{
@@ -359,10 +430,22 @@ func columnsToStruct(req *plugin.CodeGenRequest, name string, columns []pyColumn
 ifsuffix>0{
 fieldName=fmt.Sprintf("%s_%d", fieldName, suffix)
  }
-gs.Fields=append(gs.Fields, Field{
+
+f:=Field{
 Name: fieldName,
 Type: makePyType(req, c.Column),
- })
+ }
+
+ifc.embed!=nil{
+f.Type=pyType{
+InnerType: "models."+modelName(c.embed.modelType, req.Settings),
+IsArray: false,
+IsNull: false,
+ }
+f.EmbedFields=c.embed.fields
+ }
+
+gs.Fields=append(gs.Fields, f)
 seen[colName]++
  }
 return&gs
@@ -476,6 +559,7 @@ func buildQueries(conf Config, req *plugin.CodeGenRequest, structs []Struct) ([]
 columns=append(columns, pyColumn{
 id: int32(i),
 Column: c,
+embed: newGoEmbed(c.EmbedTable, structs, req.Catalog.DefaultSchema),
  })
  }
 gs=columnsToStruct(req, query.Name+"Row", columns)