Account for insert/upsert with where clause

simo7 · simo7 · commit d1a001d7ea4e · 2024-09-13T22:31:23.000+02:00
diff --git a/internal/gen.go b/internal/gen.go
@@ -889,7 +889,8 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node {
 		case ":one":
 			fetchrow := connMethodNode("fetchrow", q.ConstantName, q.ArgNodes()...)
 			f.Body = append(f.Body, assignNode("row", poet.Await(fetchrow)))
-			if strings.HasPrefix(q.ConstantName, "INSERT") || strings.HasPrefix(q.ConstantName, "UPSERT") {
+			if !strings.Contains(strings.ToUpper(q.SQL), "WHERE ") &&
+				(strings.HasPrefix(q.ConstantName, "INSERT") || strings.HasPrefix(q.ConstantName, "UPSERT")) {
 				f.Returns = q.Ret.Annotation()
 			} else {
 				f.Body = append(f.Body, poet.Node(
