Skip to content

Commit 12312df

Browse files
rayakamerayakame
authored
Feature asyncpg (#29)
* first iteration of asyncpg support * feat(sqlc-gen-better-python): Add note about Asyncpg support in README (main) * refactor: Use asyncpgWriteParams instead of aiosqliteWriteParams (main) * fragment * ci(driver/asyncpg): fix typo preventing return rows collection (#15) * chore: remove unnecessary line break in Asyncpg support note (feature/asyncpg)
1 parent 7964a2f commit 12312df

File tree

12 files changed

+288
-134
lines changed

12 files changed

+288
-134
lines changed

‎.changes/unreleased/Added-20250505-215623.yaml‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
kind: Added
2+
body: Added early driver support for `asyncpg`. Only has support for `exec`, `many` and `one`
3+
time: 2025-05-05T21:56:23.6920401+02:00
4+
custom:
5+
Author: rayakame
6+
PR: "29"

‎README.md‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,14 @@ The supported [query commands](https://docs.sqlc.dev/en/latest/reference/query-a
1313
1414
> Prepared Queries are not planned for the near future, but will be implemented sooner or later
1515
16+
> [!NOTE]
17+
> Asyncpg only has very bad support until now. It doesn't support `:execresult`, `:execrows` and `:execlastid`
18+
1619
| | `:exec` | `:execresult` | `:execrows` | `:execlastid` | `:many` | `:one` | `:copyfrom` |
17-
| --------- | ------- | ------------- | ----------- | ------------- | ------- | ------ | ----------- |
20+
| --------- |---------| ------------- | ----------- | ------------- |---------|--------| ----------- |
1821
| aiosqlite | yes | yes | yes | yes | yes | yes | no |
1922
| sqlite3 | yes | yes | yes | yes | yes | yes | no |
20-
| asyncpg | no | no | no | no | no | no | no |
23+
| asyncpg | yes | no | no | no | yes | yes | no |
2124
| psycopg2 | no | no | no | no | no | no | no |
2225
| mysql | no | no | no | no | no | no | no |
2326

‎internal/codegen/common.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@ func NewDriver(conf *core.Config) (*Driver, error) {
3333
buildPyQueryFunc = drivers.SQLite3BuildPyQueryFunc
3434
acceptedDriverCMDs = drivers.SQLite3AcceptedDriverCMDs
3535
connType = drivers.SQLite3Conn
36+
case core.SQLDriverAsyncpg:
37+
buildPyQueryFunc = drivers.AsyncpgBuildPyQueryFunc
38+
acceptedDriverCMDs = drivers.AsyncpgAcceptedDriverCMDs
39+
connType = drivers.AsyncpgConn
3640
default:
3741
return nil, fmt.Errorf("unsupported driver: %s", conf.SqlDriver.String())
3842
}

‎internal/codegen/drivers/asyncpg.go‎

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
package drivers
2+
3+
import (
4+
"fmt"
5+
"github.com/rayakame/sqlc-gen-better-python/internal/codegen/builders"
6+
"github.com/rayakame/sqlc-gen-better-python/internal/core"
7+
"github.com/sqlc-dev/plugin-sdk-go/metadata"
8+
"strconv"
9+
"strings"
10+
)
11+
12+
const AsyncpgConn = "asyncpg.Connection"
13+
14+
func AsyncpgBuildPyQueryFunc(query *core.Query, body *builders.IndentStringBuilder, args []string, retType string, isClass bool) error {
15+
indentLevel := 0
16+
params := fmt.Sprintf("conn: %s", AsyncpgConn)
17+
conn := "conn"
18+
if isClass {
19+
params = "self"
20+
conn = "self._conn"
21+
indentLevel = 1
22+
}
23+
body.WriteIndentedString(indentLevel, fmt.Sprintf("async def %s(%s", query.FuncName, params))
24+
for i, arg := range args {
25+
if i == 0 {
26+
body.WriteString(", *")
27+
}
28+
body.WriteString(fmt.Sprintf(", %s", arg))
29+
}
30+
if query.Cmd == metadata.CmdExec {
31+
body.WriteLine(fmt.Sprintf(") -> %s:", retType))
32+
body.WriteIndentedString(indentLevel+1, fmt.Sprintf("await %s.execute(%s", conn, query.ConstantName))
33+
asyncpgWriteParams(query, body)
34+
body.WriteLine(")")
35+
} else if query.Cmd == metadata.CmdOne {
36+
body.WriteLine(fmt.Sprintf(") -> typing.Optional[%s]:", retType))
37+
body.WriteIndentedString(indentLevel+1, fmt.Sprintf("row = await %s.fetchrow(%s", conn, query.ConstantName))
38+
asyncpgWriteParams(query, body)
39+
body.WriteLine(")")
40+
body.WriteIndentedLine(indentLevel+1, "if row is None:")
41+
body.WriteIndentedLine(indentLevel+2, "return None")
42+
if query.Ret.IsStruct() {
43+
body.WriteIndentedString(indentLevel+1, fmt.Sprintf("return %s(", retType))
44+
i := 0
45+
for _, col := range query.Ret.Table.Columns {
46+
if i != 0 {
47+
body.WriteString(", ")
48+
}
49+
if len(col.EmbedFields) != 0 {
50+
var inner []string
51+
body.WriteString(fmt.Sprintf("%s=%s(", col.Name, col.Type.Type))
52+
for _, embedCol := range col.EmbedFields {
53+
inner = append(inner, fmt.Sprintf("%s=row[%s]", embedCol.Name, strconv.Itoa(i)))
54+
i++
55+
}
56+
body.WriteString(strings.Join(inner, ", ") + ")")
57+
} else {
58+
body.WriteString(fmt.Sprintf("%s=row[%s]", col.Name, strconv.Itoa(i)))
59+
i++
60+
}
61+
}
62+
body.WriteLine(")")
63+
} else {
64+
body.WriteIndentedLine(indentLevel+1, fmt.Sprintf("return %s(row[0])", retType))
65+
}
66+
} else if query.Cmd == metadata.CmdMany {
67+
body.WriteLine(fmt.Sprintf(") -> typing.Sequence[%s]:", retType))
68+
body.WriteIndentedString(indentLevel+1, fmt.Sprintf("rows = await %s.fetch(%s", conn, query.ConstantName))
69+
asyncpgWriteParams(query, body)
70+
body.WriteLine(")")
71+
body.WriteIndentedLine(indentLevel+1, fmt.Sprintf("return_rows: typing.List[%s] = []", retType))
72+
body.WriteIndentedLine(indentLevel+1, "for row in rows:")
73+
if query.Ret.IsStruct() {
74+
body.WriteIndentedString(indentLevel+2, fmt.Sprintf("return_rows.append(%s(", retType))
75+
i := 0
76+
for _, col := range query.Ret.Table.Columns {
77+
if i != 0 {
78+
body.WriteString(", ")
79+
}
80+
if len(col.EmbedFields) != 0 {
81+
var inner []string
82+
body.WriteString(fmt.Sprintf("%s=%s(", col.Name, col.Type.Type))
83+
for _, embedCol := range col.EmbedFields {
84+
inner = append(inner, fmt.Sprintf("%s=row[%s]", embedCol.Name, strconv.Itoa(i)))
85+
i++
86+
}
87+
body.WriteString(strings.Join(inner, ", ") + ")")
88+
} else {
89+
body.WriteString(fmt.Sprintf("%s=row[%s]", col.Name, strconv.Itoa(i)))
90+
i++
91+
}
92+
}
93+
body.WriteLine("))")
94+
body.WriteIndentedString(indentLevel+1, "return return_rows")
95+
} else {
96+
body.WriteIndentedLine(indentLevel+2, fmt.Sprintf("return_rows.append(%s(row[0]))", retType))
97+
body.WriteIndentedString(indentLevel+1, "return return_rows")
98+
}
99+
}
100+
return nil
101+
}
102+
103+
func AsyncpgAcceptedDriverCMDs() []string {
104+
return []string{
105+
metadata.CmdExec,
106+
metadata.CmdOne,
107+
metadata.CmdMany,
108+
}
109+
}
110+
111+
func asyncpgWriteParams(query *core.Query, body *builders.IndentStringBuilder) {
112+
if len(query.Args) == 0 {
113+
return
114+
}
115+
params := ""
116+
for i, arg := range query.Args {
117+
if !arg.IsEmpty() {
118+
if i == len(query.Args)-1 {
119+
params += fmt.Sprintf(" %s", arg.Name)
120+
} else {
121+
params += fmt.Sprintf(" %s,", arg.Name)
122+
}
123+
}
124+
}
125+
body.WriteString("," + params)
126+
}

‎internal/core/enums.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ func (dr *SQLDriverType) String() string {
1111
const (
1212
SQLDriverSQLite SQLDriverType = "sqlite3"
1313
SQLDriverAioSQLite SQLDriverType = "aiosqlite"
14+
SQLDriverAsyncpg SQLDriverType = "asyncpg"
1415
)
1516

1617
const (
@@ -21,11 +22,13 @@ const (
2122
var asyncDrivers = map[SQLDriverType]bool{
2223
SQLDriverSQLite: false,
2324
SQLDriverAioSQLite: true,
25+
SQLDriverAsyncpg: true,
2426
}
2527

2628
var driversEngine = map[SQLDriverType]string{
2729
SQLDriverSQLite: "sqlite",
2830
SQLDriverAioSQLite: "sqlite",
31+
SQLDriverAsyncpg: "postgresql",
2932
}
3033

3134
var validModelTypes = map[string]struct{}{

‎pyproject.toml‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,6 @@ name = "sqlc-gen-better-python"
33
version = "0.1.0"
44
description = "Add your description here"
55
requires-python = ">=3.13"
6-
dependencies = []
6+
dependencies = [
7+
"asyncpg>=0.30.0",
8+
]

‎sqlc.yaml‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,17 @@ plugins:
33
- name: python
44
wasm:
55
url: file://sqlc-gen-better-python.wasm
6-
sha256: 9a83b2a13344cebb64f33b3deacbab37a098409bd8027b996e32f041aead9267
6+
sha256: 453e3fbe37020527d1faf55a698e3019d30d46214c92b39837140d74d164a50d
77
sql:
88
- schema: test/schema.sql
99
queries: test/queries.sql
10-
engine: sqlite
10+
engine: postgresql
1111
codegen:
1212
- out: test
1313
plugin: python
1414
options:
1515
package: test
16-
sql_driver: sqlite3
16+
sql_driver: asyncpg
1717
model_type: dataclass
1818
emit_classes: false
1919
omit_unused_structs: true

‎test/models.py‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,6 @@
55
from __future__ import annotations
66

77
__all__: typing.Sequence[str] = (
8-
"Author",
98
"Student",
109
"TestScore",
1110
)
@@ -14,13 +13,6 @@
1413
import typing
1514

1615

17-
@dataclasses.dataclass()
18-
class Author:
19-
id: int
20-
name: str
21-
bio: typing.Optional[str]
22-
23-
2416
@dataclasses.dataclass()
2517
class Student:
2618
id: int

0 commit comments

Comments
 (0)