Skip to main content

sqlc_gen_sqlx/codegen/
mod.rs

1use crate::{
2    catalog, config::Config, emit::FileEmitter, error::Error, plugin::GenerateRequestView,
3    types::TypeMap,
4};
5
6mod batch;
7mod composites;
8mod copyfrom;
9mod enums;
10mod query;
11
12pub fn generate(request: &GenerateRequestView<'_>, config: &Config) -> Result<String, Error> {
13    let mut type_map = TypeMap::new(&config.overrides, &config.copy_cheap_types);
14    let catalog_info = catalog::walk(request, &mut type_map)?;
15    let col_overrides = crate::types::build_column_overrides(&config.overrides);
16    let mut emitter = FileEmitter::new(request.sqlc_version, env!("CARGO_PKG_VERSION"));
17
18    // Emit type definitions before query code.
19    for info in &catalog_info.enums {
20        emitter.push(enums::gen_enum(info, &config.enum_derives)?);
21    }
22    for info in &catalog_info.composites {
23        emitter.push(composites::gen_composite(info, &config.composite_derives)?);
24    }
25
26    let mut module_items: Vec<proc_macro2::TokenStream> = Vec::new();
27    let mut impl_fns: Vec<proc_macro2::TokenStream> = Vec::new();
28
29    for q in request.queries.iter() {
30        let (outer, inner) = match q.cmd {
31            ":exec" => query::gen_exec(q, &type_map, config, &col_overrides)?,
32            ":execrows" => query::gen_execrows(q, &type_map, config, &col_overrides)?,
33            ":execresult" => query::gen_execresult(q, &type_map, config, &col_overrides)?,
34            ":execlastid" => query::gen_execlastid(q, &type_map, config, &col_overrides)?,
35            ":batchexec" => batch::gen_batchexec(q, &type_map, config, &col_overrides)?,
36            ":batchone" => batch::gen_batchone(q, &type_map, config, &col_overrides)?,
37            ":batchmany" => batch::gen_batchmany(q, &type_map, config, &col_overrides)?,
38            ":copyfrom" => copyfrom::gen_copyfrom(q, &type_map, config, &col_overrides)?,
39            ":one" => query::gen_one(q, &type_map, config, &col_overrides)?,
40            ":many" => query::gen_many(q, &type_map, config, &col_overrides)?,
41            cmd => {
42                eprintln!("sqlc-gen-sqlx: skipping unsupported annotation {cmd}");
43                continue;
44            }
45        };
46        module_items.push(outer);
47        impl_fns.push(inner);
48    }
49
50    for item in module_items {
51        emitter.push(item);
52    }
53
54    emitter.push(quote::quote! {
55        pub trait AsExecutor {
56            fn as_executor(&mut self) -> impl sqlx::Executor<'_, Database = sqlx::Postgres>;
57        }
58
59        impl AsExecutor for sqlx::PgPool {
60            fn as_executor(&mut self) -> impl sqlx::Executor<'_, Database = sqlx::Postgres> {
61                &*self
62            }
63        }
64
65        impl AsExecutor for &sqlx::PgPool {
66            fn as_executor(&mut self) -> impl sqlx::Executor<'_, Database = sqlx::Postgres> {
67                *self
68            }
69        }
70
71        impl AsExecutor for sqlx::PgConnection {
72            fn as_executor(&mut self) -> impl sqlx::Executor<'_, Database = sqlx::Postgres> {
73                &mut *self
74            }
75        }
76
77        impl AsExecutor for sqlx::Transaction<'_, sqlx::Postgres> {
78            fn as_executor(&mut self) -> impl sqlx::Executor<'_, Database = sqlx::Postgres> {
79                &mut **self
80            }
81        }
82
83        impl AsExecutor for sqlx::pool::PoolConnection<sqlx::Postgres> {
84            fn as_executor(&mut self) -> impl sqlx::Executor<'_, Database = sqlx::Postgres> {
85                &mut **self
86            }
87        }
88
89        impl<T: AsExecutor + ?Sized> AsExecutor for &mut T {
90            fn as_executor(&mut self) -> impl sqlx::Executor<'_, Database = sqlx::Postgres> {
91                (**self).as_executor()
92            }
93        }
94
95        pub struct Queries<E> {
96            db: E,
97        }
98
99        impl<E> Queries<E> {
100            pub fn new(db: E) -> Self {
101                Self { db }
102            }
103
104            pub fn into_inner(self) -> E {
105                self.db
106            }
107        }
108    });
109
110    if !impl_fns.is_empty() {
111        emitter.push(quote::quote! {
112            impl<E: AsExecutor> Queries<E> {
113                #(#impl_fns)*
114            }
115        });
116    }
117
118    emitter.finish()
119}