fix: enforce timeout for copyFrom

Stumble · Jul 11, 2024 · 1811445 · 1811445
1 parent 95f14a4
commit 1811445
Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go
@@ -437,6 +437,7 @@ func (i *importer) copyfromImports() fileImports {
 	})
 
 	std["context"] = struct{}{}
+	std["time"] = struct{}{}
 	if i.Settings.Go.SqlDriver == SQLDriverGoSQLDriverMySQL {
 		std["io"] = struct{}{}
 		std["fmt"] = struct{}{}

diff --git a/internal/codegen/golang/templates/wpgx/copyfromCopy.tmpl b/internal/codegen/golang/templates/wpgx/copyfromCopy.tmpl
@@ -38,10 +38,17 @@ func (r iteratorFor{{.MethodName}}) Err() error {
 {{range .Comments}}//{{.}}
 {{end -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.SlicePair}}) (int64, error) {
+{{- if gt .Option.Timeout.Milliseconds 0 }}
+    ctx, cancel := context.WithTimeout(ctx, time.Millisecond * {{.Option.Timeout.Milliseconds}})
+    defer cancel()
+{{- end}}
     q.db.CountIntent("{{.UniqueLabel}}")
 	return q.db.WCopyFrom(ctx, "{{.UniqueLabel}}", {{.TableIdentifierAsGoSlice}}, {{.Arg.ColumnNamesAsGoSlice}}, &iteratorFor{{.MethodName}}{rows: {{.Arg.Name}}})
 }
 
+// eliminate unused error
+var _ = time.Now()
+
 {{end}}
 {{end}}
 {{end}}