Skip to content

Commit 4e865d5

Browse files
asmyasnikovasmyasnikov
committed
fix codegen for database/sql
1 parent cc74fcd commit 4e865d5

1 file changed

Lines changed: 75 additions & 25 deletions

File tree

‎internal/codegen/dbsql/gen.go‎

Lines changed: 75 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -168,15 +168,22 @@ func genQueryFile(cat *pb.Catalog, pkg, sourceName string, queries []*pb.Query)
168168
var buf bytes.Buffer
169169
buf.WriteString("// Code generated by sqlc. DO NOT EDIT.\n// source: " + sourceName + "\n\n")
170170
buf.WriteString("package " + pkg + "\n\n")
171-
buf.WriteString("import (\n\t\"context\"\n\t\"database/sql\"\n)\n\n")
172-
173-
for _, q := range queries {
171+
buf.WriteString("import (\n")
172+
buf.WriteString("\t\"context\"\n\n")
173+
buf.WriteString("\t_ \"github.com/ydb-platform/ydb-go-sdk/v3\"\n")
174+
buf.WriteString("\t\"github.com/ydb-platform/ydb-go-sdk/v3/pkg/xerrors\"\n")
175+
buf.WriteString("\t\"github.com/ydb-platform/ydb-go-sdk/v3/retry\"\n")
176+
buf.WriteString(")\n\n")
177+
178+
for i, q := range queries {
174179
block, err := genQuery(cat, pkg, q)
175180
if err != nil {
176181
return nil, err
177182
}
178183
buf.Write(block)
179-
buf.WriteString("\n")
184+
if i < len(queries)-1 {
185+
buf.WriteString("\n")
186+
}
180187
}
181188
return buf.Bytes(), nil
182189
}
@@ -261,17 +268,27 @@ func genOne(cat *pb.Catalog, name, constName string, q *pb.Query) string {
261268
sig = "ctx context.Context, arg " + name + "Params"
262269
args = argList(params)
263270
}
264-
b.WriteString("func (q *Queries) " + name + "(" + sig + ") (" + retType + ", error) {\n")
265-
b.WriteString("\trow := q.db.QueryRowContext(ctx, " + constName + ", " + args + ")\n")
266-
b.WriteString("\tvar i " + retType + "\n")
267-
b.WriteString("\terr := row.Scan(" + scanTargets(cols) + ")\n")
268-
b.WriteString("\tif err != nil {\n\t\treturn i, err\n\t}\n")
269-
b.WriteString("\treturn i, nil\n}\n\n")
271+
b.WriteString("func (q *Queries) " + name + "(" + sig + ") (*" + retType + ", error) {\n")
272+
b.WriteString("\ti, err := retry.RetryWithResult(ctx, func(ctx context.Context) (*" + retType + ", error) {\n")
273+
b.WriteString("\t\trow := q.db.QueryRowContext(ctx, " + constName + ", " + args + ")\n")
274+
b.WriteString("\t\tvar i " + retType + "\n")
275+
b.WriteString("\t\terr := row.Scan(" + scanTargets(cols) + ")\n")
276+
b.WriteString("\t\tif err != nil {\n")
277+
b.WriteString("\t\t\treturn nil, xerrors.WithStackTrace(err)\n")
278+
b.WriteString("\t\t}\n\n")
279+
b.WriteString("\t\treturn &i, nil\n")
280+
b.WriteString("\t}, retry.WithLabel(\"" + name + "\"))\n")
281+
b.WriteString("\tif err != nil {\n")
282+
b.WriteString("\t\treturn nil, xerrors.WithStackTrace(err)\n")
283+
b.WriteString("\t}\n\n")
284+
b.WriteString("\treturn i, nil\n")
285+
b.WriteString("}\n")
270286
return b.String()
271287
}
272288

273289
func genMany(cat *pb.Catalog, name, constName string, q *pb.Query) string {
274290
var b bytes.Buffer
291+
params := q.GetParams()
275292
cols := q.GetColumns()
276293
rowType := "interface{}"
277294
if len(cols) > 0 {
@@ -285,18 +302,41 @@ func genMany(cat *pb.Catalog, name, constName string, q *pb.Query) string {
285302
rowType = name + "Row"
286303
}
287304
}
288-
b.WriteString("func (q *Queries) " + name + "(ctx context.Context) ([]" + rowType + ", error) {\n")
289-
b.WriteString("\trows, err := q.db.QueryContext(ctx, " + constName + ")\n")
290-
b.WriteString("\tif err != nil {\n\t\treturn nil, err\n\t}\n")
291-
b.WriteString("\tdefer rows.Close()\n")
292-
b.WriteString("\tvar items []" + rowType + "\n")
293-
b.WriteString("\tfor rows.Next() {\n")
294-
b.WriteString("\t\tvar i " + rowType + "\n")
295-
b.WriteString("\t\tif err := rows.Scan(" + scanTargets(cols) + "); err != nil {\n\t\t\treturn nil, err\n\t\t}\n")
296-
b.WriteString("\t\titems = append(items, i)\n")
297-
b.WriteString("\t}\n")
298-
b.WriteString("\tif err := rows.Err(); err != nil {\n\t\treturn nil, err\n\t}\n")
299-
b.WriteString("\treturn items, nil\n}\n\n")
305+
sig := "ctx context.Context"
306+
args := ""
307+
if len(params) > 0 {
308+
sig = "ctx context.Context, arg " + name + "Params"
309+
args = argList(params)
310+
}
311+
b.WriteString("func (q *Queries) " + name + "(" + sig + ") ([]" + rowType + ", error) {\n")
312+
b.WriteString("\titems, err := retry.RetryWithResult(ctx, func(ctx context.Context) ([]" + rowType + ", error) {\n")
313+
if len(params) > 0 {
314+
b.WriteString("\t\trows, err := q.db.QueryContext(ctx, " + constName + ", " + args + ")\n")
315+
} else {
316+
b.WriteString("\t\trows, err := q.db.QueryContext(ctx, " + constName + ")\n")
317+
}
318+
b.WriteString("\t\tif err != nil {\n")
319+
b.WriteString("\t\t\treturn nil, err\n")
320+
b.WriteString("\t\t}\n")
321+
b.WriteString("\t\tdefer rows.Close()\n")
322+
b.WriteString("\t\tvar items []" + rowType + "\n")
323+
b.WriteString("\t\tfor rows.Next() {\n")
324+
b.WriteString("\t\t\tvar i " + rowType + "\n")
325+
b.WriteString("\t\t\tif err := rows.Scan(" + scanTargets(cols) + "); err != nil {\n")
326+
b.WriteString("\t\t\t\treturn nil, xerrors.WithStackTrace(err)\n")
327+
b.WriteString("\t\t\t}\n")
328+
b.WriteString("\t\t\titems = append(items, i)\n")
329+
b.WriteString("\t\t}\n")
330+
b.WriteString("\t\tif err := rows.Err(); err != nil {\n")
331+
b.WriteString("\t\t\treturn nil, xerrors.WithStackTrace(err)\n")
332+
b.WriteString("\t\t}\n\n")
333+
b.WriteString("\t\treturn items, nil\n")
334+
b.WriteString("\t}, retry.WithLabel(\"" + name + "\"))\n")
335+
b.WriteString("\tif err != nil {\n")
336+
b.WriteString("\t\treturn nil, xerrors.WithStackTrace(err)\n")
337+
b.WriteString("\t}\n\n")
338+
b.WriteString("\treturn items, nil\n")
339+
b.WriteString("}\n")
300340
return b.String()
301341
}
302342

@@ -310,12 +350,22 @@ func genExec(name, constName string, q *pb.Query) string {
310350
args = argList(params)
311351
}
312352
b.WriteString("func (q *Queries) " + name + "(" + sig + ") error {\n")
353+
b.WriteString("\terr := retry.Retry(ctx, func(ctx context.Context) error {\n")
313354
if len(params) > 0 {
314-
b.WriteString("\t_, err := q.db.ExecContext(ctx, " + constName + ", " + args + ")\n")
355+
b.WriteString("\t\t_, err := q.db.ExecContext(ctx, " + constName + ", " + args + ")\n")
315356
} else {
316-
b.WriteString("\t_, err := q.db.ExecContext(ctx, " + constName + ")\n")
357+
b.WriteString("\t\t_, err := q.db.ExecContext(ctx, " + constName + ")\n")
317358
}
318-
b.WriteString("\treturn err\n}\n\n")
359+
b.WriteString("\t\tif err != nil {\n")
360+
b.WriteString("\t\t\treturn xerrors.WithStackTrace(err)\n")
361+
b.WriteString("\t\t}\n\n")
362+
b.WriteString("\t\treturn nil\n")
363+
b.WriteString("\t}, retry.WithLabel(\"" + name + "\"))\n")
364+
b.WriteString("\tif err != nil {\n")
365+
b.WriteString("\t\treturn xerrors.WithStackTrace(err)\n")
366+
b.WriteString("\t}\n\n")
367+
b.WriteString("\treturn nil\n")
368+
b.WriteString("}\n")
319369
return b.String()
320370
}
321371

0 commit comments

Comments
 (0)