Fix :copyfrom by adding columns kwarg

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

@@ -0,0 +1,6 @@
+kind: Changed
+body: Improved `:execrows` performance for `asyncpg` and added speedup option for `:execrows`
+time: 2025-05-31T00:08:42.9297814+02:00
+custom:
+    Author: rayakame
+    PR: "112"

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

@@ -0,0 +1,6 @@
+kind: Fixed
+body: Added `columns` kwarg to `:copyfrom` for `asyncpg` to fix inserts for columns with default values
+time: 2025-05-31T00:07:49.4105598+02:00
+custom:
+    Author: rayakame
+    PR: "112"

internal/codegen/common.go

@@ -7,7 +7,7 @@ import (
 	"github.com/rayakame/sqlc-gen-better-python/internal/core"
 )
 
-type TypeBuildPyQueryFunc func(*core.Query, *builders.IndentStringBuilder, []core.FunctionArg, core.PyType, bool) error
+type TypeBuildPyQueryFunc func(*core.Query, *builders.IndentStringBuilder, []core.FunctionArg, core.PyType, *core.Config) error
 type TypeAcceptedDriverCMDs func() []string
 type TypeDriverTypeCheckingHook func() []string
 type TypeDriverBuildQueryResults func(*builders.IndentStringBuilder) string

internal/codegen/drivers/aiosqlite.go

@@ -159,13 +159,13 @@ func AiosqliteBuildQueryResults(body *builders.IndentStringBuilder) string {
 	return "QueryResults"
 }
 
-func AioSQLiteBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, isClass bool) error {
+func AioSQLiteBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, conf *core.Config) error {
 	indentLevel := 0
 	params := fmt.Sprintf("conn: %s", AioSQLiteConn)
 	conn := "conn"
 	asyncFunc := "async "
 	docstringConnType := AioSQLiteConn
-	if isClass {
+	if conf.EmitClasses {
 		params = "self"
 		conn = "self._conn"
 		indentLevel = 1

internal/codegen/drivers/asyncpg.go

@@ -47,13 +47,13 @@ func AsyncpgBuildQueryResults(body *builders.IndentStringBuilder) string {
 	return "QueryResults"
 }
 
-func AsyncpgBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, isClass bool) error {
+func AsyncpgBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, conf *core.Config) error {
 	indentLevel := 0
 	params := fmt.Sprintf("conn: %s", AsyncpgConn)
 	conn := "conn"
 	asyncFunc := "async "
 	docstringConnType := AsyncpgConn
-	if isClass {
+	if conf.EmitClasses {
 		params = "self"
 		conn = "self._conn"
 		indentLevel = 1
@@ -84,27 +84,38 @@ func AsyncpgBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuild
 	} else if query.Cmd == metadata.CmdExecRows {
 		body.WriteLine(fmt.Sprintf(") -> %s:", retType.Type))
 		body.WriteQueryFunctionDocstring(indentLevel+1, query, docstringConnType, args, retType)
-		body.WriteIndentedString(indentLevel+1, fmt.Sprintf("result = await %s.execute(%s", conn, query.ConstantName))
+		body.WriteIndentedString(indentLevel+1, fmt.Sprintf("r = await %s.execute(%s", conn, query.ConstantName))
 		asyncpgWriteParams(query, body)
 		body.WriteLine(")")
-		body.WriteIndentedLine(indentLevel+1, "return int(result.split()[-1]) if result.split()[-1].isdigit() else 0")
+		if conf.Speedups {
+			body.WriteIndentedLine(indentLevel+1, "return int(n) if (n := r.split()[-1]).isdigit() else 0")
+		} else {
+			body.WriteIndentedLine(indentLevel+1, "return int(n) if (p := r.split()) and (n := p[-1]).isdigit() else 0")
+		}
 	} else if query.Cmd == metadata.CmdCopyFrom {
 		body.WriteLine(fmt.Sprintf(") -> %s:", retType.Type))
 		body.WriteQueryFunctionDocstring(indentLevel+1, query, docstringConnType, args, retType)
 		body.WriteIndentedLine(indentLevel+1, "records = [")
 		params := ""
+		columns := ``
 		for i, arg := range query.Args[0].Table.Columns {
 			if i == len(query.Args[0].Table.Columns)-1 && i != 0 {
 				params += fmt.Sprintf("%s.%s", "param", arg.Name)
+				columns += fmt.Sprintf(`"%s"`, arg.Name)
 			} else {
 				params += fmt.Sprintf("%s.%s, ", "param", arg.Name)
+				columns += fmt.Sprintf(`"%s", `, arg.Name)
 			}
 		}
 		body.WriteIndentedLine(indentLevel+2, fmt.Sprintf("(%s)", params))
 		body.WriteIndentedLine(indentLevel+2, fmt.Sprintf("for param in %s", query.Args[0].Name))
 		body.WriteIndentedLine(indentLevel+1, "]")
-		body.WriteIndentedLine(indentLevel+1, fmt.Sprintf(`result = await %s.copy_records_to_table("%s", records=records)`, conn, query.Table.Name))
-		body.WriteIndentedLine(indentLevel+1, "return int(result.split()[-1]) if result.split()[-1].isdigit() else 0")
+		body.WriteIndentedLine(indentLevel+1, fmt.Sprintf(`r = await %s.copy_records_to_table("%s", columns=[%s], records=records)`, conn, query.Table.Name, columns))
+		if conf.Speedups {
+			body.WriteIndentedLine(indentLevel+1, "return int(n) if (n := r.split()[-1]).isdigit() else 0")
+		} else {
+			body.WriteIndentedLine(indentLevel+1, "return int(n) if (p := r.split()) and (n := p[-1]).isdigit() else 0")
+		}
 	} else if query.Cmd == metadata.CmdOne {
 		body.WriteLine(fmt.Sprintf(") -> %s | None:", retType.Type))
 		body.WriteQueryFunctionDocstring(indentLevel+1, query, docstringConnType, args, retType)

internal/codegen/drivers/sqlite3.go

@@ -160,12 +160,12 @@ func SQLite3BuildQueryResults(body *builders.IndentStringBuilder) string {
 	return "QueryResults"
 }
 
-func SQLite3BuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, isClass bool) error {
+func SQLite3BuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []core.FunctionArg, retType core.PyType, conf *core.Config) error {
 	indentLevel := 0
 	params := fmt.Sprintf("conn: %s", SQLite3Conn)
 	conn := "conn"
 	docstringConnType := SQLite3Conn
-	if isClass {
+	if conf.EmitClasses {
 		params = "self"
 		conn = "self._conn"
 		indentLevel = 1

internal/codegen/queries.go

@@ -155,7 +155,7 @@ func (dr *Driver) buildPyQueriesFile(imp *core.Importer, queries []core.Query, s
 			Type:    retType,
 		}
 		allNames = append(allNames, addedPyTableNames...)
-		err := dr.buildPyQueryFunc(&query, funcBody, args, returnType, dr.conf.EmitClasses)
+		err := dr.buildPyQueryFunc(&query, funcBody, args, returnType, dr.conf)
 		if err != nil {
 			return nil, err
 		}

sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
   - name: python
     wasm:
       url: file://sqlc-gen-better-python.wasm
-      sha256: 93df043f217f5bc2b9ce70c1283d36676b9ffe8f7ab9b5e8379db20db045e92f
+      sha256: 750bbaa7a45acdce18c4c06ea4384c4df7e5f3e0e51ed6c8badcaea2dbe6d14d
 sql:
   - schema: test/schema.sql
     queries: test/queries.sql

test/driver_aiosqlite/attrs/classes/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_aiosqlite/attrs/classes/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations
 

test/driver_aiosqlite/attrs/classes/queries.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing queries from file queries.sql."""
 from __future__ import annotations
 

test/driver_aiosqlite/attrs/functions/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_aiosqlite/attrs/functions/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations
 

test/driver_aiosqlite/attrs/functions/queries.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing queries from file queries.sql."""
 from __future__ import annotations
 

test/driver_aiosqlite/dataclass/classes/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_aiosqlite/dataclass/classes/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations
 

test/driver_aiosqlite/dataclass/classes/queries.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing queries from file queries.sql."""
 from __future__ import annotations
 

test/driver_aiosqlite/dataclass/functions/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_aiosqlite/dataclass/functions/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations
 

test/driver_aiosqlite/dataclass/functions/queries.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing queries from file queries.sql."""
 from __future__ import annotations
 

test/driver_aiosqlite/msgspec/classes/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_aiosqlite/msgspec/classes/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations
 

test/driver_aiosqlite/msgspec/classes/queries.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing queries from file queries.sql."""
 from __future__ import annotations
 

test/driver_aiosqlite/msgspec/functions/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_aiosqlite/msgspec/functions/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations
 

test/driver_aiosqlite/msgspec/functions/queries.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing queries from file queries.sql."""
 from __future__ import annotations
 

test/driver_aiosqlite/sqlc.yaml

@@ -3,7 +3,7 @@ plugins:
   - name: python
     wasm:
       url: file://sqlc-gen-better-python.wasm
-      sha256: 93df043f217f5bc2b9ce70c1283d36676b9ffe8f7ab9b5e8379db20db045e92f
+      sha256: 750bbaa7a45acdce18c4c06ea4384c4df7e5f3e0e51ed6c8badcaea2dbe6d14d
 sql:
   - schema: schema.sql
     queries: queries.sql

test/driver_asyncpg/attrs/classes/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_asyncpg/attrs/classes/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations
 

test/driver_asyncpg/attrs/classes/queries.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing queries from file queries.sql."""
 from __future__ import annotations
 
@@ -944,8 +944,8 @@ async def create_rows_one_test_postgres_type(self, *, id_: int, serial_test: int
             The number of affected rows. This will be 0 for queries like `CREATE TABLE`.
 
         """
-        result = await self._conn.execute(CREATE_ROWS_ONE_TEST_POSTGRES_TYPE, id_, serial_test, serial4_test, bigserial_test, smallserial_test, int_test, bigint_test, smallint_test, float_test, double_precision_test, real_test, numeric_test, money_test, bool_test, json_test, jsonb_test, bytea_test, date_test, time_test, timetz_test, timestamp_test, timestamptz_test, interval_test, text_test, varchar_test, bpchar_test, char_test, citext_test, uuid_test, inet_test, cidr_test, macaddr_test, macaddr8_test, ltree_test, lquery_test, ltxtquery_test)
-        return int(result.split()[-1]) if result.split()[-1].isdigit() else 0
+        r = await self._conn.execute(CREATE_ROWS_ONE_TEST_POSTGRES_TYPE, id_, serial_test, serial4_test, bigserial_test, smallserial_test, int_test, bigint_test, smallint_test, float_test, double_precision_test, real_test, numeric_test, money_test, bool_test, json_test, jsonb_test, bytea_test, date_test, time_test, timetz_test, timestamp_test, timestamptz_test, interval_test, text_test, varchar_test, bpchar_test, char_test, citext_test, uuid_test, inet_test, cidr_test, macaddr_test, macaddr8_test, ltree_test, lquery_test, ltxtquery_test)
+        return int(n) if (n := r.split()[-1]).isdigit() else 0
 
     async def create_rows_table(self) -> int:
         """Execute SQL query with `name: CreateRowsTable :execrows` and return the number of affected rows.
@@ -964,8 +964,8 @@ async def create_rows_table(self) -> int:
             The number of affected rows. This will be 0 for queries like `CREATE TABLE`.
 
         """
-        result = await self._conn.execute(CREATE_ROWS_TABLE)
-        return int(result.split()[-1]) if result.split()[-1].isdigit() else 0
+        r = await self._conn.execute(CREATE_ROWS_TABLE)
+        return int(n) if (n := r.split()[-1]).isdigit() else 0
 
     async def delete_one_result_test_postgres_type(self, *, id_: int) -> str:
         """Execute and return the result of SQL query with `name: DeleteOneResultTestPostgresType :execresult`.
@@ -1007,8 +1007,8 @@ async def delete_one_rows_test_postgres_type(self, *, id_: int) -> int:
             The number of affected rows. This will be 0 for queries like `CREATE TABLE`.
 
         """
-        result = await self._conn.execute(DELETE_ONE_ROWS_TEST_POSTGRES_TYPE, id_)
-        return int(result.split()[-1]) if result.split()[-1].isdigit() else 0
+        r = await self._conn.execute(DELETE_ONE_ROWS_TEST_POSTGRES_TYPE, id_)
+        return int(n) if (n := r.split()[-1]).isdigit() else 0
 
     async def delete_one_test_postgres_inner_type(self, *, table_id: int) -> None:
         """Execute SQL query with `name: DeleteOneTestPostgresInnerType :exec`.
@@ -1413,8 +1413,8 @@ async def test_copy_from(self, *, params: collections.abc.Sequence[TestCopyFromP
             (param.id, param.float_test, param.int_test)
             for param in params
         ]
-        result = await self._conn.copy_records_to_table("test_copy_from", records=records)
-        return int(result.split()[-1]) if result.split()[-1].isdigit() else 0
+        r = await self._conn.copy_records_to_table("test_copy_from", columns=["id", "float_test", "int_test"], records=records)
+        return int(n) if (n := r.split()[-1]).isdigit() else 0
 
     async def update_result_test_postgres_type(self, *, id_: int) -> str:
         """Execute and return the result of SQL query with `name: UpdateResultTestPostgresType :execresult`.
@@ -1456,5 +1456,5 @@ async def update_rows_test_postgres_type(self, *, id_: int) -> int:
             The number of affected rows. This will be 0 for queries like `CREATE TABLE`.
 
         """
-        result = await self._conn.execute(UPDATE_ROWS_TEST_POSTGRES_TYPE, id_)
-        return int(result.split()[-1]) if result.split()[-1].isdigit() else 0
+        r = await self._conn.execute(UPDATE_ROWS_TEST_POSTGRES_TYPE, id_)
+        return int(n) if (n := r.split()[-1]).isdigit() else 0

test/driver_asyncpg/attrs/functions/__init__.py

@@ -1,5 +1,5 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Package containing queries and models automatically generated using sqlc-gen-better-python."""

test/driver_asyncpg/attrs/functions/models.py

@@ -1,7 +1,7 @@
 # Code generated by sqlc. DO NOT EDIT.
 # versions:
 #   sqlc v1.28.0
-#   sqlc-gen-better-python v0.4.2
+#   sqlc-gen-better-python v0.4.3
 """Module containing models."""
 from __future__ import annotations