Skip to main content

sqlcx_core/generator/python/
asyncpg.rs

1// asyncpg driver. Emits queries.py only. Native $N placeholders, positional
2// args, async functions against asyncpg.Connection. Rows hydrate via
3// dict(row) because asyncpg Records are mapping-like.
4
5use crate::error::Result;
6use crate::generator::python::common::{
7    PyBodyCtx, PyDriverShape, PyTypeMap, generate_driver_files,
8};
9use crate::generator::{DriverGenerator, GeneratedFile};
10use crate::ir::{QueryCommand, QueryDef, SqlcxIR};
11
12pub struct AsyncpgGenerator;
13
14impl PyTypeMap for AsyncpgGenerator {}
15
16impl PyDriverShape for AsyncpgGenerator {
17    fn driver_import(&self) -> &'static str {
18        "from asyncpg import Connection"
19    }
20    fn connection_type(&self) -> &'static str {
21        "Connection"
22    }
23    fn is_async(&self) -> bool {
24        true
25    }
26    fn rewrite_sql(&self, query: &QueryDef) -> String {
27        // asyncpg uses native $N — no rewrite.
28        query.sql.clone()
29    }
30    fn build_params_arg(&self, query: &QueryDef) -> String {
31        if query.params.is_empty() {
32            return String::new();
33        }
34        // Positional args appended after the SQL const: ", params.a, params.b"
35        // in $N order (params already sorted by index).
36        let args: Vec<String> = query
37            .params
38            .iter()
39            .map(|p| format!("params.{}", p.name))
40            .collect();
41        format!(", {}", args.join(", "))
42    }
43    fn render_body(&self, ctx: &PyBodyCtx<'_>) -> (String, String) {
44        let (sc, rt, pa) = (ctx.sql_const, ctx.row_type, ctx.params_arg);
45        match ctx.command {
46            QueryCommand::One => (
47                format!("{rt} | None"),
48                format!(
49                    "    row = await conn.fetchrow({sc}{pa})\n    if row is None:\n        return None\n    return {rt}(**dict(row))"
50                ),
51            ),
52            QueryCommand::Many => (
53                format!("list[{rt}]"),
54                format!(
55                    "    rows = await conn.fetch({sc}{pa})\n    return [{rt}(**dict(row)) for row in rows]"
56                ),
57            ),
58            QueryCommand::Exec => (
59                "None".to_string(),
60                format!("    await conn.execute({sc}{pa})"),
61            ),
62            QueryCommand::ExecResult => (
63                "int".to_string(),
64                format!(
65                    "    result = await conn.execute({sc}{pa})\n    return int(result.split()[-1])"
66                ),
67            ),
68        }
69    }
70}
71
72impl DriverGenerator for AsyncpgGenerator {
73    fn generate(&self, ir: &SqlcxIR) -> Result<Vec<GeneratedFile>> {
74        generate_driver_files(self, ir)
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use crate::generator::python::common::generate_queries_file;
82    use crate::parser::DatabaseParser;
83    use crate::parser::postgres::PostgresParser;
84
85    fn parse_fixture_ir() -> SqlcxIR {
86        let schema_sql = include_str!("../../../../../tests/fixtures/schema.sql");
87        let queries_sql = include_str!("../../../../../tests/fixtures/queries/users.sql");
88        let parser = PostgresParser::new();
89        let (tables, enums) = parser.parse_schema(schema_sql).unwrap();
90        let queries = parser
91            .parse_queries(queries_sql, &tables, &enums, "queries/users.sql")
92            .unwrap();
93        SqlcxIR {
94            tables,
95            queries,
96            enums,
97        }
98    }
99
100    #[test]
101    fn generates_asyncpg_query_functions() {
102        let ir = parse_fixture_ir();
103        let content = generate_queries_file(&AsyncpgGenerator, &ir.queries);
104        assert!(content.contains("from asyncpg import Connection"));
105        assert!(content.contains("async def get_user"));
106        assert!(content.contains("await conn.fetchrow"));
107        insta::assert_snapshot!("asyncpg_queries", content);
108    }
109}