11package golang
22
33import (
4- "bufio"
5- "bytes"
64 "context"
75 "errors"
86 "fmt"
9- "go/format"
107 "strings"
11- "text/template"
128
139 "github.com/sqlc-dev/sqlc/internal/codegen/golang/opts"
14- "github.com/sqlc-dev/sqlc/internal/codegen/sdk"
1510 "github.com/sqlc-dev/sqlc/internal/metadata"
1611 "github.com/sqlc-dev/sqlc/internal/plugin"
1712)
@@ -171,7 +166,7 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
171166 Structs : structs ,
172167 }
173168
174- tctx := tmplCtx {
169+ tctx := & tmplCtx {
175170 EmitInterface : options .EmitInterface ,
176171 EmitJSONTags : options .EmitJsonTags ,
177172 JsonTagsIDUppercase : options .JsonTagsIdUppercase ,
@@ -209,64 +204,9 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
209204 return nil , errors .New (":batch* commands are only supported by pgx" )
210205 }
211206
212- funcMap := template.FuncMap {
213- "lowerTitle" : sdk .LowerTitle ,
214- "comment" : sdk .DoubleSlashComment ,
215- "escape" : sdk .EscapeBacktick ,
216- "imports" : i .Imports ,
217- "hasImports" : i .HasImports ,
218- "hasPrefix" : strings .HasPrefix ,
219-
220- // These methods are Go specific, they do not belong in the codegen package
221- // (as that is language independent)
222- "dbarg" : tctx .codegenDbarg ,
223- "emitPreparedQueries" : tctx .codegenEmitPreparedQueries ,
224- "queryMethod" : tctx .codegenQueryMethod ,
225- "queryRetval" : tctx .codegenQueryRetval ,
226- }
227-
228- tmpl := template .Must (
229- template .New ("table" ).
230- Funcs (funcMap ).
231- ParseFS (
232- templates ,
233- "templates/*.tmpl" ,
234- "templates/*/*.tmpl" ,
235- ),
236- )
237-
238207 output := map [string ]string {}
239208
240- execute := func (name , templateName string ) error {
241- imports := i .Imports (name )
242- replacedQueries := replaceConflictedArg (imports , queries )
243-
244- var b bytes.Buffer
245- w := bufio .NewWriter (& b )
246- tctx .SourceName = name
247- tctx .GoQueries = replacedQueries
248- err := tmpl .ExecuteTemplate (w , templateName , & tctx )
249- w .Flush ()
250- if err != nil {
251- return err
252- }
253- code , err := format .Source (b .Bytes ())
254- if err != nil {
255- fmt .Println (b .String ())
256- return fmt .Errorf ("source error: %w" , err )
257- }
258-
259- if templateName == "queryFile" && options .OutputFilesSuffix != "" {
260- name += options .OutputFilesSuffix
261- }
262-
263- if ! strings .HasSuffix (name , ".go" ) {
264- name += ".go"
265- }
266- output [name ] = string (code )
267- return nil
268- }
269-
209+ // File names
270210 dbFileName := "db.go"
271211 if options .OutputDbFileName != "" {
272212 dbFileName = options .OutputDbFileName
@@ -283,46 +223,89 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
283223 if options .OutputCopyfromFileName != "" {
284224 copyfromFileName = options .OutputCopyfromFileName
285225 }
286-
287226 batchFileName := "batch.go"
288227 if options .OutputBatchFileName != "" {
289228 batchFileName = options .OutputBatchFileName
290229 }
291230
292- if err := execute (dbFileName , "dbFile" ); err != nil {
293- return nil , err
231+ // Generate db.go
232+ tctx .SourceName = dbFileName
233+ tctx .GoQueries = replaceConflictedArg (i .Imports (dbFileName ), queries )
234+ gen := NewCodeGenerator (tctx , i )
235+
236+ code , err := gen .GenerateDBFile ()
237+ if err != nil {
238+ return nil , fmt .Errorf ("db file error: %w" , err )
294239 }
295- if err := execute (modelsFileName , "modelsFile" ); err != nil {
296- return nil , err
240+ output [dbFileName ] = string (code )
241+
242+ // Generate models.go
243+ tctx .SourceName = modelsFileName
244+ tctx .GoQueries = replaceConflictedArg (i .Imports (modelsFileName ), queries )
245+ code , err = gen .GenerateModelsFile ()
246+ if err != nil {
247+ return nil , fmt .Errorf ("models file error: %w" , err )
297248 }
249+ output [modelsFileName ] = string (code )
250+
251+ // Generate querier.go
298252 if options .EmitInterface {
299- if err := execute (querierFileName , "interfaceFile" ); err != nil {
300- return nil , err
253+ tctx .SourceName = querierFileName
254+ tctx .GoQueries = replaceConflictedArg (i .Imports (querierFileName ), queries )
255+ code , err = gen .GenerateQuerierFile ()
256+ if err != nil {
257+ return nil , fmt .Errorf ("querier file error: %w" , err )
301258 }
259+ output [querierFileName ] = string (code )
302260 }
261+
262+ // Generate copyfrom.go
303263 if tctx .UsesCopyFrom {
304- if err := execute (copyfromFileName , "copyfromFile" ); err != nil {
305- return nil , err
264+ tctx .SourceName = copyfromFileName
265+ tctx .GoQueries = replaceConflictedArg (i .Imports (copyfromFileName ), queries )
266+ code , err = gen .GenerateCopyFromFile ()
267+ if err != nil {
268+ return nil , fmt .Errorf ("copyfrom file error: %w" , err )
306269 }
270+ output [copyfromFileName ] = string (code )
307271 }
272+
273+ // Generate batch.go
308274 if tctx .UsesBatch {
309- if err := execute (batchFileName , "batchFile" ); err != nil {
310- return nil , err
275+ tctx .SourceName = batchFileName
276+ tctx .GoQueries = replaceConflictedArg (i .Imports (batchFileName ), queries )
277+ code , err = gen .GenerateBatchFile ()
278+ if err != nil {
279+ return nil , fmt .Errorf ("batch file error: %w" , err )
311280 }
281+ output [batchFileName ] = string (code )
312282 }
313283
314- files := map [string ]struct {}{}
284+ // Generate query files
285+ sourceFiles := map [string ]struct {}{}
315286 for _ , gq := range queries {
316- files [gq .SourceName ] = struct {}{}
287+ sourceFiles [gq .SourceName ] = struct {}{}
317288 }
318289
319- for source := range files {
320- if err := execute (source , "queryFile" ); err != nil {
321- return nil , err
290+ for source := range sourceFiles {
291+ tctx .SourceName = source
292+ tctx .GoQueries = replaceConflictedArg (i .Imports (source ), queries )
293+ code , err = gen .GenerateQueryFile (source )
294+ if err != nil {
295+ return nil , fmt .Errorf ("query file error for %s: %w" , source , err )
322296 }
297+
298+ filename := source
299+ if options .OutputFilesSuffix != "" {
300+ filename += options .OutputFilesSuffix
301+ }
302+ if ! strings .HasSuffix (filename , ".go" ) {
303+ filename += ".go"
304+ }
305+ output [filename ] = string (code )
323306 }
324- resp := plugin.GenerateResponse {}
325307
308+ resp := plugin.GenerateResponse {}
326309 for filename , code := range output {
327310 resp .Files = append (resp .Files , & plugin.File {
328311 Name : filename ,
0 commit comments