Skip to content

Commit 7d9f9ec

Browse files
rayakamerayakame
authored
Fix :copyfrom by adding columns kwarg (#112)
* pipeline * fix(asyncpg): Fix an issue with columns argument missing in copy_records_to_table function call (bugfix/copyfrom-colums-kwarg) * feat(asyncpg): Add speedup logic for copy_records_to_table method - Updated AsyncpgBuildPyQueryFunc to include speedup logic for copy_records_to_table method - Now, the method returns 'int(n)' if 'n' is a digit, else returns 0 - Implemented the same logic for both 'execute' and 'copy_records_to_table' - Issue: bugfix/copyfrom-colums-kwarg * fragment
1 parent 24859f0 commit 7d9f9ec

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

69 files changed

+158
-132
lines changed

‎.changes/unreleased/Changed-20250531-000842.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
kind: Changed
2+
body: Improved `:execrows` performance for `asyncpg` and added speedup option for `:execrows`
3+
time: 2025-05-31T00:08:42.9297814+02:00
4+
custom:
5+
Author: rayakame
6+
PR: "112"

‎.changes/unreleased/Fixed-20250531-000749.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
kind: Fixed
2+
body: Added `columns` kwarg to `:copyfrom` for `asyncpg` to fix inserts for columns with default values
3+
time: 2025-05-31T00:07:49.4105598+02:00
4+
custom:
5+
Author: rayakame
6+
PR: "112"

‎internal/codegen/common.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import (
77
"github.com/rayakame/sqlc-gen-better-python/internal/core"
88
)
99

10-
type TypeBuildPyQueryFunc func(*core.Query, *builders.IndentStringBuilder, []core.FunctionArg, core.PyType, bool) error
10+
type TypeBuildPyQueryFunc func(*core.Query, *builders.IndentStringBuilder, []core.FunctionArg, core.PyType, *core.Config) error
1111
type TypeAcceptedDriverCMDs func() []string
1212
type TypeDriverTypeCheckingHook func() []string
1313
type TypeDriverBuildQueryResults func(*builders.IndentStringBuilder) string

‎internal/codegen/drivers/aiosqlite.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -159,13 +159,13 @@ func AiosqliteBuildQueryResults(body *builders.IndentStringBuilder) string {
159159
return "QueryResults"
160160
}
161161

162-
func AioSQLiteBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, isClass bool) error {
162+
func AioSQLiteBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, conf *core.Config) error {
163163
indentLevel := 0
164164
params := fmt.Sprintf("conn: %s", AioSQLiteConn)
165165
conn := "conn"
166166
asyncFunc := "async "
167167
docstringConnType := AioSQLiteConn
168-
if isClass {
168+
if conf.EmitClasses {
169169
params = "self"
170170
conn = "self._conn"
171171
indentLevel = 1

‎internal/codegen/drivers/asyncpg.go‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -47,13 +47,13 @@ func AsyncpgBuildQueryResults(body *builders.IndentStringBuilder) string {
4747
return "QueryResults"
4848
}
4949

50-
func AsyncpgBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, isClass bool) error {
50+
func AsyncpgBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, conf *core.Config) error {
5151
indentLevel := 0
5252
params := fmt.Sprintf("conn: %s", AsyncpgConn)
5353
conn := "conn"
5454
asyncFunc := "async "
5555
docstringConnType := AsyncpgConn
56-
if isClass {
56+
if conf.EmitClasses {
5757
params = "self"
5858
conn = "self._conn"
5959
indentLevel = 1
@@ -84,27 +84,38 @@ func AsyncpgBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuild
8484
} else if query.Cmd == metadata.CmdExecRows {
8585
body.WriteLine(fmt.Sprintf(") -> %s:", retType.Type))
8686
body.WriteQueryFunctionDocstring(indentLevel+1, query, docstringConnType, args, retType)
87-
body.WriteIndentedString(indentLevel+1, fmt.Sprintf("result = await %s.execute(%s", conn, query.ConstantName))
87+
body.WriteIndentedString(indentLevel+1, fmt.Sprintf("r = await %s.execute(%s", conn, query.ConstantName))
8888
asyncpgWriteParams(query, body)
8989
body.WriteLine(")")
90-
body.WriteIndentedLine(indentLevel+1, "return int(result.split()[-1]) if result.split()[-1].isdigit() else 0")
90+
if conf.Speedups {
91+
body.WriteIndentedLine(indentLevel+1, "return int(n) if (n := r.split()[-1]).isdigit() else 0")
92+
} else {
93+
body.WriteIndentedLine(indentLevel+1, "return int(n) if (p := r.split()) and (n := p[-1]).isdigit() else 0")
94+
}
9195
} else if query.Cmd == metadata.CmdCopyFrom {
9296
body.WriteLine(fmt.Sprintf(") -> %s:", retType.Type))
9397
body.WriteQueryFunctionDocstring(indentLevel+1, query, docstringConnType, args, retType)
9498
body.WriteIndentedLine(indentLevel+1, "records = [")
9599
params := ""
100+
columns := ``
96101
for i, arg := range query.Args[0].Table.Columns {
97102
if i == len(query.Args[0].Table.Columns)-1 && i != 0 {
98103
params += fmt.Sprintf("%s.%s", "param", arg.Name)
104+
columns += fmt.Sprintf(`"%s"`, arg.Name)
99105
} else {
100106
params += fmt.Sprintf("%s.%s, ", "param", arg.Name)
107+
columns += fmt.Sprintf(`"%s", `, arg.Name)
101108
}
102109
}
103110
body.WriteIndentedLine(indentLevel+2, fmt.Sprintf("(%s)", params))
104111
body.WriteIndentedLine(indentLevel+2, fmt.Sprintf("for param in %s", query.Args[0].Name))
105112
body.WriteIndentedLine(indentLevel+1, "]")
106-
body.WriteIndentedLine(indentLevel+1, fmt.Sprintf(`result = await %s.copy_records_to_table("%s", records=records)`, conn, query.Table.Name))
107-
body.WriteIndentedLine(indentLevel+1, "return int(result.split()[-1]) if result.split()[-1].isdigit() else 0")
113+
body.WriteIndentedLine(indentLevel+1, fmt.Sprintf(`r = await %s.copy_records_to_table("%s", columns=[%s], records=records)`, conn, query.Table.Name, columns))
114+
if conf.Speedups {
115+
body.WriteIndentedLine(indentLevel+1, "return int(n) if (n := r.split()[-1]).isdigit() else 0")
116+
} else {
117+
body.WriteIndentedLine(indentLevel+1, "return int(n) if (p := r.split()) and (n := p[-1]).isdigit() else 0")
118+
}
108119
} else if query.Cmd == metadata.CmdOne {
109120
body.WriteLine(fmt.Sprintf(") -> %s | None:", retType.Type))
110121
body.WriteQueryFunctionDocstring(indentLevel+1, query, docstringConnType, args, retType)

‎internal/codegen/drivers/sqlite3.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -160,12 +160,12 @@ func SQLite3BuildQueryResults(body *builders.IndentStringBuilder) string {
160160
return "QueryResults"
161161
}
162162

163-
func SQLite3BuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, isClass bool) error {
163+
func SQLite3BuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, conf *core.Config) error {
164164
indentLevel := 0
165165
params := fmt.Sprintf("conn: %s", SQLite3Conn)
166166
conn := "conn"
167167
docstringConnType := SQLite3Conn
168-
if isClass {
168+
if conf.EmitClasses {
169169
params = "self"
170170
conn = "self._conn"
171171
indentLevel = 1

‎internal/codegen/queries.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,7 @@ func (dr *Driver) buildPyQueriesFile(imp *core.Importer, queries []core.Query, s
155155
Type: retType,
156156
}
157157
allNames = append(allNames, addedPyTableNames...)
158-
err := dr.buildPyQueryFunc(&query, funcBody, args, returnType, dr.conf.EmitClasses)
158+
err := dr.buildPyQueryFunc(&query, funcBody, args, returnType, dr.conf)
159159
if err != nil {
160160
return nil, err
161161
}

‎sqlc.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ plugins:
33
- name: python
44
wasm:
55
url: file://sqlc-gen-better-python.wasm
6-
sha256: 93df043f217f5bc2b9ce70c1283d36676b9ffe8f7ab9b5e8379db20db045e92f
6+
sha256: 750bbaa7a45acdce18c4c06ea4384c4df7e5f3e0e51ed6c8badcaea2dbe6d14d
77
sql:
88
- schema: test/schema.sql
99
queries: test/queries.sql

‎test/driver_aiosqlite/attrs/classes/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
# Code generated by sqlc. DO NOT EDIT.
22
# versions:
33
# sqlc v1.28.0
4-
# sqlc-gen-better-python v0.4.2
4+
# sqlc-gen-better-python v0.4.3
55
"""Package containing queries and models automatically generated using sqlc-gen-better-python."""

‎test/driver_aiosqlite/attrs/classes/models.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
# Code generated by sqlc. DO NOT EDIT.
22
# versions:
33
# sqlc v1.28.0
4-
# sqlc-gen-better-python v0.4.2
4+
# sqlc-gen-better-python v0.4.3
55
"""Module containing models."""
66
from __future__ import annotations
77

0 commit comments

Comments
 (0)