Implement exec rows and exec result

sqlc-dev · Dec 15, 2021 · cc8cf38 · cc8cf38
1 parent 963cf0d
commit cc8cf38
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 0 deletions.
diff --git a/internal/codegen/python/gen.go b/internal/codegen/python/gen.go
@@ -1003,8 +1003,15 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node {
 				f.Body = append(f.Body, exec)
 				f.Returns = poet.Constant(nil)
 			case ":execrows":
+				f.Body = append(f.Body,
+					assignNode("result", exec),
+					poet.Return(poet.Attribute(poet.Name("result"), "rowcount")),
+				)
 				f.Returns = poet.Name("int")
 			case ":execresult":
+				f.Body = append(f.Body,
+					poet.Return(exec),
+				)
 				f.Returns = typeRefNode("sqlalchemy", "engine", "Result")
 			default:
 				panic("unknown cmd " + q.Cmd)
@@ -1089,8 +1096,15 @@ func buildQueryTree(ctx *pyTmplCtx, i *importer, source string) *pyast.Node {
 				f.Body = append(f.Body, poet.Await(exec))
 				f.Returns = poet.Constant(nil)
 			case ":execrows":
+				f.Body = append(f.Body,
+					assignNode("result", poet.Await(exec)),
+					poet.Return(poet.Attribute(poet.Name("result"), "rowcount")),
+				)
 				f.Returns = poet.Name("int")
 			case ":execresult":
+				f.Body = append(f.Body,
+					poet.Return(poet.Await(exec)),
+				)
 				f.Returns = typeRefNode("sqlalchemy", "engine", "Result")
 			default:
 				panic("unknown cmd " + q.Cmd)

diff --git a/internal/endtoend/testdata/exec_result/python_postgresql/python/models.py b/internal/endtoend/testdata/exec_result/python_postgresql/python/models.py
@@ -1,7 +1,13 @@
 # Code generated by sqlc. DO NOT EDIT.
+
+
 import dataclasses
 
 
+
+
 @dataclasses.dataclass()
 class Bar:
     id: int
+
+
diff --git a/internal/endtoend/testdata/exec_rows/python_postgresql/python/models.py b/internal/endtoend/testdata/exec_rows/python_postgresql/python/models.py
@@ -1,7 +1,13 @@
 # Code generated by sqlc. DO NOT EDIT.
+
+
 import dataclasses
 
 
+
+
 @dataclasses.dataclass()
 class Bar:
     id: int
+
+