Skip to content

Commit 643c37c

Browse files
committed
Add support for :copyfrom in Python code generation
1 parent b67b396 commit 643c37c

File tree

3 files changed

+30
-15
lines changed

3 files changed

+30
-15
lines changed

internal/endtoend/testdata/copyfrom/python/query.py

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,11 @@
2020
"""
2121

2222

23+
CREATE_AUTHORS_NAMED = """-- name: create_authors_named \\:copyfrom
24+
INSERT INTO authors (name, bio) VALUES (:p1, :p2)
25+
"""
26+
27+
2328
class Querier:
2429
def __init__(self, conn: sqlalchemy.engine.Connection):
2530
self._conn = conn
@@ -35,7 +40,11 @@ def create_author(self, *, name: str, bio: str) -> Optional[models.Author]:
3540
)
3641

3742
def create_authors(self, arg_list: List[Any]) -> int:
38-
result = self._conn.execute(sqlalchemy.text(CREATE_AUTHORS), arg_list)
43+
result = self._conn.executemany(sqlalchemy.text(CREATE_AUTHORS), arg_list)
44+
return result.rowcount
45+
46+
def create_authors_named(self, arg_list: List[Any]) -> int:
47+
result = self._conn.executemany(sqlalchemy.text(CREATE_AUTHORS_NAMED), arg_list)
3948
return result.rowcount
4049

4150

@@ -54,5 +63,9 @@ async def create_author(self, *, name: str, bio: str) -> Optional[models.Author]
5463
)
5564

5665
async def create_authors(self, arg_list: List[Any]) -> int:
57-
result = await self._conn.execute(sqlalchemy.text(CREATE_AUTHORS), arg_list)
66+
result = await self._conn.executemany(sqlalchemy.text(CREATE_AUTHORS), arg_list)
67+
return result.rowcount
68+
69+
async def create_authors_named(self, arg_list: List[Any]) -> int:
70+
result = await self._conn.executemany(sqlalchemy.text(CREATE_AUTHORS_NAMED), arg_list)
5871
return result.rowcount

internal/endtoend/testdata/copyfrom/sqlc.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ plugins:
33
- name: py
44
wasm:
55
url: file://../../../../bin/sqlc-gen-python.wasm
6-
sha256: "e9601f06d9e7d20e2908cf83b8e06688c426b148aead0ebeb282776cfb020b6b"
6+
sha256: "32036421ce8137b6b1314a86069a04100fafc925627f132f38588e5b35ce8b54"
77
sql:
88
- schema: "schema.sql"
99
queries: "query.sql"

internal/gen.go

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -211,11 +211,10 @@ func (q Query) ArgDictNode() *pyast.Node {
211211
}
212212

213213
// BuildCopyFromBody generates the method body for :copyfrom commands.
214-
//
215-
// IMPORTANT: This implementation uses executemany for batch inserts, which is
216-
// faster than individual inserts but not as fast as PostgreSQL's native COPY.
217-
// For maximum performance with PostgreSQL, consider using psycopg's copy_from
218-
// directly via conn.connection.driver_connection.
214+
// This implementation explicitly uses executemany() for batch inserts, which is
215+
// significantly faster than individual inserts. For PostgreSQL specifically,
216+
// using the native COPY protocol would be even faster, but executemany provides
217+
// good performance while maintaining cross-database compatibility.
219218
func (q Query) BuildCopyFromBody(isAsync bool) []*pyast.Node {
220219
var body []*pyast.Node
221220

@@ -238,15 +237,18 @@ func (q Query) BuildCopyFromBody(isAsync bool) []*pyast.Node {
238237
dataVar = argName
239238
}
240239

241-
// Build the execute call with the SQL and parameter list
242-
// SQLAlchemy detects the list and uses executemany internally
240+
// Use executemany explicitly for batch inserts
241+
// Build the SQL text object first
242+
sqlText := poet.Node(&pyast.Call{
243+
Func: poet.Attribute(poet.Name("sqlalchemy"), "text"),
244+
Args: []*pyast.Node{poet.Name(q.ConstantName)},
245+
})
246+
247+
// Call executemany with the SQL and parameter list
243248
execCall := poet.Node(&pyast.Call{
244-
Func: poet.Attribute(poet.Name("self._conn"), "execute"),
249+
Func: poet.Attribute(poet.Name("self._conn"), "executemany"),
245250
Args: []*pyast.Node{
246-
poet.Node(&pyast.Call{
247-
Func: poet.Attribute(poet.Name("sqlalchemy"), "text"),
248-
Args: []*pyast.Node{poet.Name(q.ConstantName)},
249-
}),
251+
sqlText,
250252
poet.Name(dataVar),
251253
},
252254
})

0 commit comments

Comments
 (0)