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            type Exec<'c>: sqlx::Executor<'c, Database = sqlx::Postgres>
57            where
58                Self: 'c;
59            fn as_executor(&mut self) -> Self::Exec<'_>;
60        }
61
62        impl AsExecutor for sqlx::PgPool {
63            type Exec<'c> = &'c sqlx::PgPool where Self: 'c;
64            fn as_executor(&mut self) -> Self::Exec<'_> { &*self }
65        }
66
67        impl AsExecutor for &sqlx::PgPool {
68            type Exec<'c> = &'c sqlx::PgPool where Self: 'c;
69            fn as_executor(&mut self) -> Self::Exec<'_> { *self }
70        }
71
72        impl AsExecutor for sqlx::PgConnection {
73            type Exec<'c> = &'c mut sqlx::PgConnection where Self: 'c;
74            fn as_executor(&mut self) -> Self::Exec<'_> { self }
75        }
76
77        impl AsExecutor for sqlx::Transaction<'_, sqlx::Postgres> {
78            type Exec<'c> = &'c mut sqlx::PgConnection where Self: 'c;
79            fn as_executor(&mut self) -> Self::Exec<'_> { &mut **self }
80        }
81
82        impl AsExecutor for sqlx::pool::PoolConnection<sqlx::Postgres> {
83            type Exec<'c> = &'c mut sqlx::PgConnection where Self: 'c;
84            fn as_executor(&mut self) -> Self::Exec<'_> { &mut **self }
85        }
86
87        impl<T: AsExecutor + ?Sized> AsExecutor for &mut T {
88            type Exec<'c> = T::Exec<'c> where Self: 'c;
89            fn as_executor(&mut self) -> Self::Exec<'_> { (**self).as_executor() }
90        }
91
92        pub struct Queries<E> {
93            db: E,
94        }
95
96        impl<E> Queries<E> {
97            pub fn new(db: E) -> Self {
98                Self { db }
99            }
100
101            pub fn into_inner(self) -> E {
102                self.db
103            }
104        }
105    });
106
107    if !impl_fns.is_empty() {
108        emitter.push(quote::quote! {
109            impl<E: AsExecutor> Queries<E> {
110                #(#impl_fns)*
111            }
112        });
113    }
114
115    emitter.finish()
116}