sqlcx_core/generator/python/
asyncpg.rs1use 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 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 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}