Skip to main content

qusql_sqlx_type_macro/
lib.rs

1//! Implementation of typed db query macros
2//!
3//! Used the exposed macros from the qusql-sqlx-type crate
4//!
5#![forbid(unsafe_code)]
6
7use std::path::PathBuf;
8use std::sync::{Arc, Mutex};
9use std::time::SystemTime;
10
11use ariadne::{Color, Label, Report, ReportKind, Source};
12use once_cell::sync::Lazy;
13use proc_macro::TokenStream;
14use proc_macro2::Span;
15use quote::{format_ident, quote, quote_spanned};
16use qusql_type::schema::{parse_schemas, Schemas};
17use qusql_type::{
18    type_statement, ByteToChar, Issue, SQLArguments, SQLDialect, SelectTypeColumn, TypeOptions,
19};
20use syn::spanned::Spanned;
21use syn::{parse::Parse, punctuated::Punctuated, Expr, Ident, LitStr, Token};
22use yoke::{Yoke, Yokeable};
23
24/// Path where the sqlx-type-schema.sql file can be found
25static SCHEMA_PATH: Lazy<PathBuf> = Lazy::new(|| {
26    let mut schema_path: PathBuf = std::env::var("CARGO_MANIFEST_DIR")
27        .expect("`CARGO_schema_path` must be set")
28        .into();
29
30    schema_path.push("sqlx-type-schema.sql");
31
32    if !schema_path.exists() {
33        use serde::Deserialize;
34        use std::process::Command;
35
36        let cargo = std::env::var("CARGO").expect("`CARGO` must be set");
37        schema_path.pop();
38
39        let output = Command::new(cargo)
40            .args(["metadata", "--format-version=1"])
41            .current_dir(&schema_path)
42            .env_remove("__CARGO_FIX_PLZ")
43            .output()
44            .expect("Could not fetch metadata");
45
46        /// Struct used to deserialize the cargo betadata
47        #[derive(Deserialize)]
48        struct CargoMetadata {
49            /// The root of the workspace
50            workspace_root: PathBuf,
51        }
52
53        let metadata: CargoMetadata =
54            serde_json::from_slice(&output.stdout).expect("Invalid `cargo metadata` output");
55
56        schema_path = metadata.workspace_root;
57        schema_path.push("sqlx-type-schema.sql");
58    }
59    if !schema_path.exists() {
60        panic!("Unable to locate sqlx-type-schema.sql");
61    }
62    schema_path
63});
64
65/// Construct a none color report for an issue
66fn issue_to_report(issue: Issue, b2c: &ByteToChar) -> Report<'static, std::ops::Range<usize>> {
67    let span = b2c.map_span(issue.span);
68    let mut builder = Report::build(
69        match issue.level {
70            qusql_type::Level::Warning => ReportKind::Warning,
71            qusql_type::Level::Error => ReportKind::Error,
72        },
73        span.clone(),
74    )
75    .with_config(ariadne::Config::default().with_color(false))
76    .with_label(
77        Label::new(span)
78            .with_order(-1)
79            .with_priority(-1)
80            .with_message(issue.message),
81    );
82    for frag in issue.fragments {
83        builder =
84            builder.with_label(Label::new(b2c.map_span(frag.span)).with_message(frag.message));
85    }
86    builder.finish()
87}
88
89/// Construct a color report for an issue
90fn issue_to_report_color(
91    issue: Issue,
92    b2c: &ByteToChar,
93) -> Report<'static, std::ops::Range<usize>> {
94    let span = b2c.map_span(issue.span);
95    let mut builder = Report::build(
96        match issue.level {
97            qusql_type::Level::Warning => ReportKind::Warning,
98            qusql_type::Level::Error => ReportKind::Error,
99        },
100        span.clone(),
101    )
102    .with_label(
103        Label::new(span)
104            .with_color(match issue.level {
105                qusql_type::Level::Warning => Color::Yellow,
106                qusql_type::Level::Error => Color::Red,
107            })
108            .with_order(-1)
109            .with_priority(-1)
110            .with_message(issue.message),
111    );
112    for frag in issue.fragments {
113        builder = builder.with_label(
114            Label::new(b2c.map_span(frag.span))
115                .with_color(Color::Blue)
116                .with_message(frag.message),
117        );
118    }
119    builder.finish()
120}
121
122/// Source name adaptor for ariadne
123struct NamedSource<'a>(&'a str, Source<&'a str>);
124
125impl<'a> ariadne::Cache<()> for &NamedSource<'a> {
126    type Storage = &'a str;
127
128    fn display<'b>(&self, _: &'b ()) -> Option<impl std::fmt::Display + 'b> {
129        Some(self.0.to_string())
130    }
131
132    fn fetch(&mut self, _: &()) -> Result<&Source<Self::Storage>, impl std::fmt::Debug> {
133        Ok::<_, ()>(&self.1)
134    }
135}
136
137#[derive(Yokeable)]
138/// `Yokeable` wrapper that lets `Schemas<'a>` be stored in a `Yoke`.
139struct SchemasYoke<'a>(Schemas<'a>);
140
141/// Parsed schema and its source string, kept alive together so neither
142/// outlives the other.  Stored behind an `Arc` so callers can cheaply
143/// clone a handle without holding the mutex.
144struct SchemaCacheEntry {
145    /// The parsed schema yoked to the source it borrows from.
146    schemas: Yoke<SchemasYoke<'static>, String>,
147    /// SQL dialect detected from the schema file header.
148    dialect: SQLDialect,
149    /// File byte-length at the time of parsing, used for cache invalidation.
150    file_len: u64,
151    /// File modification time at the time of parsing, used for cache invalidation.
152    modified: SystemTime,
153    /// Hash of the schema content
154    hash: u64,
155}
156
157/// Global cache of the most recently parsed schema, invalidated by file mtime/size.
158static SCHEMA_CACHE: Mutex<Option<Arc<SchemaCacheEntry>>> = Mutex::new(None);
159
160/// Return the parsed schema, reparsing from disk only when the file's mtime or
161/// size has changed.  A single `stat()` is performed on every macro invocation;
162/// the expensive parse happens only on the first call and after each on-disk
163/// change.  The source string and the `Schemas` that borrows from it are kept
164/// alive together inside a `Yoke` - no memory is leaked.
165fn get_schemas() -> Arc<SchemaCacheEntry> {
166    let path = SCHEMA_PATH.as_path();
167    let meta = std::fs::metadata(path).unwrap_or_else(|e| panic!("Cannot stat {path:?}: {e}"));
168    let file_len = meta.len();
169    let modified = meta.modified().unwrap_or(SystemTime::UNIX_EPOCH);
170
171    let mut cache = SCHEMA_CACHE.lock().expect("schema cache lock poisoned");
172    if let Some(entry) = cache.as_ref() {
173        if entry.file_len == file_len && entry.modified == modified {
174            return Arc::clone(entry);
175        }
176    }
177
178    // Cache miss or stale - read and reparse.
179    let src_string = std::fs::read_to_string(path)
180        .unwrap_or_else(|e| panic!("Unable to read schema from {path:?}: {e}"));
181
182    // Detect dialect from the first two lines before moving src_box into the
183    // Yoke cart.  The result is moved into the closure by value.
184    let dialect = {
185        let header = src_string.lines().take(2).collect::<Vec<_>>().join(" ");
186        if header.contains("qusql-type-variant: postgis") || header.contains("sql-product: postgis")
187        {
188            SQLDialect::PostGIS
189        } else if header.contains("sql-product: postgres") {
190            SQLDialect::PostgreSQL
191        } else if header.contains("sql-product: sqlite") {
192            SQLDialect::Sqlite
193        } else {
194            SQLDialect::MariaDB
195        }
196    };
197
198    let dialect_for_closure = dialect.clone();
199
200    // Compute the schema hash from the fresh source before we move it into the cart.
201    let schema_hash = {
202        use std::hash::{DefaultHasher, Hash, Hasher};
203        let mut hasher = DefaultHasher::new();
204        src_string.hash(&mut hasher);
205        hasher.finish()
206    };
207
208    let schemas = Yoke::attach_to_cart(src_string, move |src| {
209        let options = TypeOptions::new().dialect(dialect_for_closure);
210        let mut issues = qusql_type::Issues::new(src);
211        let parsed = parse_schemas(src, &mut issues, &options);
212        if !issues.is_ok() {
213            let b2c = ByteToChar::new(src.as_bytes());
214            let source = NamedSource("sqlx-type-schema.sql", Source::from(src));
215            let mut err = false;
216            for issue in issues.into_vec() {
217                if issue.level == qusql_type::Level::Error {
218                    err = true;
219                }
220                issue_to_report_color(issue, &b2c).eprint(&source).unwrap();
221            }
222            if err {
223                panic!("Errors processing sqlx-type-schema.sql");
224            }
225        }
226        SchemasYoke(parsed)
227    });
228
229    let entry = Arc::new(SchemaCacheEntry {
230        schemas,
231        dialect,
232        file_len,
233        modified,
234        hash: schema_hash,
235    });
236    *cache = Some(Arc::clone(&entry));
237    entry
238}
239
240/// Produce quoted arguments for a query
241fn quote_args(
242    errors: &mut Vec<proc_macro2::TokenStream>,
243    query: &str,
244    last_span: Span,
245    args: &[Expr],
246    arguments: &[(qusql_type::ArgumentKey<'_>, qusql_type::FullType)],
247    dialect: &SQLDialect,
248) -> (proc_macro2::TokenStream, proc_macro2::TokenStream) {
249    let cls = match dialect {
250        SQLDialect::MariaDB => quote!(sqlx::mysql::MySql),
251        SQLDialect::Sqlite => quote!(sqlx::sqlite::Sqlite),
252        SQLDialect::PostgreSQL | SQLDialect::PostGIS => quote!(sqlx::postgres::Postgres),
253    };
254
255    let mut at = Vec::new();
256    let inv = qusql_type::FullType::invalid();
257    for (k, v) in arguments {
258        match k {
259            qusql_type::ArgumentKey::Index(i) => {
260                while at.len() <= *i {
261                    at.push(&inv);
262                }
263                at[*i] = v;
264            }
265            qusql_type::ArgumentKey::Identifier(_) => {
266                errors.push(
267                    syn::Error::new(last_span.span(), "Named arguments not supported")
268                        .to_compile_error(),
269                );
270            }
271        }
272    }
273
274    if at.len() > args.len() {
275        errors.push(
276            syn::Error::new(
277                last_span,
278                format!("Expected {} additional arguments", at.len() - args.len()),
279            )
280            .to_compile_error(),
281        );
282    }
283
284    if let Some(args) = args.get(at.len()..) {
285        for arg in args {
286            errors.push(syn::Error::new(arg.span(), "unexpected argument").to_compile_error());
287        }
288    }
289
290    let arg_names = (0..args.len())
291        .map(|i| format_ident!("arg{}", i))
292        .collect::<Vec<_>>();
293
294    let mut arg_bindings = Vec::new();
295    let mut arg_add = Vec::new();
296
297    let mut list_lengths = Vec::new();
298
299    for ((qa, ta), name) in args.iter().zip(at).zip(&arg_names) {
300        let mut t = match ta.t {
301            qusql_type::Type::U8 => quote! {u8},
302            qusql_type::Type::I8 => quote! {i8},
303            qusql_type::Type::U16 => quote! {u16},
304            qusql_type::Type::I16 => quote! {i16},
305            qusql_type::Type::U24 => quote! {u32},
306            qusql_type::Type::I24 => quote! {i32},
307            qusql_type::Type::U32 => quote! {u32},
308            qusql_type::Type::I32 => quote! {i32},
309            qusql_type::Type::U64 => quote! {u64},
310            qusql_type::Type::I64 => quote! {i64},
311            qusql_type::Type::Base(qusql_type::BaseType::Any) => quote! {qusql_sqlx_type::Any},
312            qusql_type::Type::Base(qusql_type::BaseType::Bool) => quote! {bool},
313            qusql_type::Type::Base(qusql_type::BaseType::Bytes) => quote! {&[u8]},
314            qusql_type::Type::Base(qusql_type::BaseType::Date) => quote! {qusql_sqlx_type::Date},
315            qusql_type::Type::Base(qusql_type::BaseType::DateTime) => {
316                quote! {qusql_sqlx_type::DateTime}
317            }
318            qusql_type::Type::Base(qusql_type::BaseType::Float) => quote! {qusql_sqlx_type::Float},
319            qusql_type::Type::Base(qusql_type::BaseType::Integer) => {
320                quote! {qusql_sqlx_type::Integer}
321            }
322            qusql_type::Type::Base(qusql_type::BaseType::String) => quote! {&str},
323            qusql_type::Type::Base(qusql_type::BaseType::Time) => quote! {qusql_sqlx_type::Time},
324            qusql_type::Type::Base(qusql_type::BaseType::TimeInterval) => todo!("time_interval"),
325            qusql_type::Type::Base(qusql_type::BaseType::TimeStamp) => {
326                quote! {qusql_sqlx_type::Timestamp}
327            }
328            qusql_type::Type::Base(qusql_type::BaseType::Uuid) => quote! {qusql_sqlx_type::Uuid},
329            qusql_type::Type::Null => todo!("null"),
330            qusql_type::Type::Invalid => quote! {std::convert::Infallible},
331            qusql_type::Type::Enum(_) => quote! {&str},
332            qusql_type::Type::Set(_) => quote! {&str},
333            qusql_type::Type::Args(_, _) => todo!("args"),
334            qusql_type::Type::F32 => quote! {f32},
335            qusql_type::Type::F64 => quote! {f64},
336            qusql_type::Type::JSON => quote! {qusql_sqlx_type::Any},
337            qusql_type::Type::Geometry => quote! {qusql_sqlx_type::Any},
338            qusql_type::Type::Range(_) => quote! {qusql_sqlx_type::Any},
339            qusql_type::Type::Array(_) => quote! {qusql_sqlx_type::Any},
340        };
341        if !ta.not_null {
342            t = quote! {Option<#t>}
343        }
344        let span = qa.span();
345        if ta.list_hack {
346            list_lengths.push(quote!(#name.len()));
347            arg_bindings.push(quote_spanned! {span=>
348                let #name = &(#qa);
349                args_count += #name.len();
350                for v in #name.iter() {
351                    size_hints += ::sqlx::encode::Encode::<#cls>::size_hint(v);
352                }
353                if false {
354                    qusql_sqlx_type::check_arg_list_hack::<#t, _>(#name);
355                    ::std::panic!();
356                }
357            });
358            arg_add.push(quote!(
359                for v in #name.iter() {
360                    e = e.and_then(|()| query_args.add(v));
361                }
362            ));
363        } else {
364            arg_bindings.push(quote_spanned! {span=>
365                let #name = &(#qa);
366                args_count += 1;
367                size_hints += ::sqlx::encode::Encode::<#cls>::size_hint(#name);
368                if false {
369                    qusql_sqlx_type::check_arg::<#t, _>(#name);
370                    ::std::panic!();
371                }
372            });
373            arg_add.push(quote!(e = e.and_then(|()| query_args.add(#name));));
374        }
375    }
376
377    let query = if list_lengths.is_empty() {
378        quote!(#query)
379    } else {
380        quote!(
381            &qusql_sqlx_type::convert_list_query(#query, &[#(#list_lengths),*])
382        )
383    };
384
385    (
386        quote! {
387            let mut size_hints = 0;
388            let mut args_count = 0;
389            #(#arg_bindings)*
390
391            let mut query_args = <#cls as ::sqlx::database::Database>::Arguments::default();
392            query_args.reserve(args_count, size_hints);
393            let mut e = Ok(());
394            #(#arg_add)*
395            let query_args = e.and_then(|()| Ok(query_args));
396        },
397        query,
398    )
399}
400
401/// Output an [Issue] as a compile error
402fn issues_to_errors(issues: Vec<Issue>, source: &str, span: Span) -> Vec<proc_macro2::TokenStream> {
403    if !issues.is_empty() {
404        let b2c = ByteToChar::new(source.as_bytes());
405        let source = NamedSource("", Source::from(source));
406        let mut err = false;
407        let mut out = Vec::new();
408        for issue in issues {
409            if issue.level == qusql_type::Level::Error {
410                err = true;
411            }
412            let r = issue_to_report(issue, &b2c);
413            r.write(&source, &mut out).unwrap();
414        }
415        if err {
416            return vec![syn::Error::new(span, String::from_utf8(out).unwrap()).to_compile_error()];
417        }
418    }
419    Vec::new()
420}
421
422/// Construct row struct members, and fill in statements
423fn construct_row(
424    columns: &[SelectTypeColumn],
425) -> (Vec<proc_macro2::TokenStream>, Vec<proc_macro2::TokenStream>) {
426    let mut row_members = Vec::new();
427    let mut row_construct = Vec::new();
428    for (i, c) in columns.iter().enumerate() {
429        let mut t = match c.type_.t {
430            qusql_type::Type::U8 => quote! {u8},
431            qusql_type::Type::I8 => quote! {i8},
432            qusql_type::Type::U16 => quote! {u16},
433            qusql_type::Type::I16 => quote! {i16},
434            qusql_type::Type::U24 => quote! {u32},
435            qusql_type::Type::I24 => quote! {i32},
436            qusql_type::Type::U32 => quote! {u32},
437            qusql_type::Type::I32 => quote! {i32},
438            qusql_type::Type::U64 => quote! {u64},
439            qusql_type::Type::I64 => quote! {i64},
440            qusql_type::Type::Base(qusql_type::BaseType::Any) => todo!("from_any"),
441            qusql_type::Type::Base(qusql_type::BaseType::Bool) => quote! {bool},
442            qusql_type::Type::Base(qusql_type::BaseType::Bytes) => quote! {Vec<u8>},
443            qusql_type::Type::Base(qusql_type::BaseType::Date) => quote! {chrono::NaiveDate},
444            qusql_type::Type::Base(qusql_type::BaseType::DateTime) => {
445                quote! {chrono::NaiveDateTime}
446            }
447            qusql_type::Type::Base(qusql_type::BaseType::Float) => quote! {f64},
448            qusql_type::Type::Base(qusql_type::BaseType::Integer) => quote! {i64},
449            qusql_type::Type::Base(qusql_type::BaseType::String) => quote! {String},
450            qusql_type::Type::Base(qusql_type::BaseType::Time) => todo!("from_time"),
451            qusql_type::Type::Base(qusql_type::BaseType::TimeInterval) => {
452                todo!("from_time_interval")
453            }
454            qusql_type::Type::Base(qusql_type::BaseType::TimeStamp) => {
455                quote! {sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>}
456            }
457            qusql_type::Type::Base(qusql_type::BaseType::Uuid) => {
458                quote! {qusql_sqlx_type::UuidValue}
459            }
460            qusql_type::Type::Null => todo!("from_null"),
461            qusql_type::Type::Invalid => quote! {i64},
462            qusql_type::Type::Enum(_) => quote! {String},
463            qusql_type::Type::Set(_) => quote! {String},
464            qusql_type::Type::Args(_, _) => todo!("from_args"),
465            qusql_type::Type::F32 => quote! {f32},
466            qusql_type::Type::F64 => quote! {f64},
467            qusql_type::Type::JSON => quote! {String},
468            qusql_type::Type::Geometry => quote! {Vec<u8>},
469            qusql_type::Type::Range(_) => quote! {Vec<u8>},
470            qusql_type::Type::Array(_) => quote! {qusql_sqlx_type::Any},
471        };
472        let name = match &c.name {
473            Some(v) => v,
474            None => continue,
475        };
476
477        // Handle sqlx "column!" convention: strip trailing ! and force not_null
478        let name_str = name.value;
479        let (name_str, force_not_null) = if let Some(stripped) = name_str.strip_suffix('!') {
480            (stripped, true)
481        } else {
482            (name_str, false)
483        };
484
485        let ident = String::from("r#") + name_str;
486        let ident: Ident = if let Ok(ident) = syn::parse_str(&ident) {
487            ident
488        } else {
489            // TODO error
490            //errors.push(syn::Error::new(span, String::from_utf8(out).unwrap()).to_compile_error().into());
491            continue;
492        };
493
494        if !c.type_.not_null && !force_not_null {
495            t = quote! {Option<#t>};
496        }
497        row_members.push(quote! {
498            #ident : #t
499        });
500        row_construct.push(quote! {
501            #ident: sqlx::Row::get(&row, #i)
502        });
503    }
504    (row_members, row_construct)
505}
506
507/// Parsed query! macro
508struct Query {
509    /// The query expression
510    query: String,
511    /// The span of the query expression
512    query_span: Span,
513    /// The arguments to supply to the query
514    args: Vec<Expr>,
515    /// The last span parsed
516    last_span: Span,
517}
518
519impl Parse for Query {
520    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
521        let query_ = Punctuated::<LitStr, Token![+]>::parse_separated_nonempty(input)?;
522        let query: String = query_.iter().map(LitStr::value).collect();
523        let query_span = query_.span();
524        let mut last_span = query_span;
525        let mut args = Vec::new();
526        while !input.is_empty() {
527            let _ = input.parse::<syn::token::Comma>()?;
528            if input.is_empty() {
529                break;
530            }
531            let arg = input.parse::<Expr>()?;
532            last_span = arg.span();
533            args.push(arg);
534        }
535        Ok(Self {
536            query,
537            query_span,
538            args,
539            last_span,
540        })
541    }
542}
543
544/// Statically checked SQL query, similarly to sqlx::query!.
545///
546/// This expands to an instance of query::Map that outputs an ad-hoc anonymous struct type.
547#[proc_macro]
548pub fn query(input: TokenStream) -> TokenStream {
549    let query = syn::parse_macro_input!(input as Query);
550    let cache = get_schemas();
551    let (schemas, dialect, schema_hash) = (cache.schemas.get(), &cache.dialect, cache.hash);
552    let options = TypeOptions::new()
553        .dialect(dialect.clone())
554        .arguments(match &dialect {
555            SQLDialect::MariaDB => SQLArguments::QuestionMark,
556            SQLDialect::Sqlite => SQLArguments::QuestionMark,
557            SQLDialect::PostgreSQL | SQLDialect::PostGIS => SQLArguments::Dollar,
558        })
559        .list_hack(true);
560    let mut issues = qusql_type::Issues::new(&query.query);
561    let stmt = type_statement(&schemas.0, &query.query, &mut issues, &options);
562    let sp = SCHEMA_PATH.as_path().to_str().unwrap();
563    let mut errors = issues_to_errors(issues.into_vec(), &query.query, query.query_span);
564    match &stmt {
565        qusql_type::StatementType::Select { columns, arguments } => {
566            let (args_tokens, q) = quote_args(
567                &mut errors,
568                &query.query,
569                query.last_span,
570                &query.args,
571                arguments,
572                dialect,
573            );
574            let (row_members, row_construct) = construct_row(columns);
575            let s = quote! { {
576                use ::sqlx::Arguments as _;
577                let _ = std::include_bytes!(#sp);
578                #(#errors; )*
579                #args_tokens;
580
581                struct Row {
582                    #(#row_members),*
583                };
584                sqlx::__query_with_result(#q, query_args).map(|row|
585                    Row{
586                        #(#row_construct),*
587                    }
588                )
589            }};
590            s.into()
591        }
592        qusql_type::StatementType::Delete {
593            arguments,
594            returning,
595        } => {
596            let (args_tokens, q) = quote_args(
597                &mut errors,
598                &query.query,
599                query.last_span,
600                &query.args,
601                arguments,
602                dialect,
603            );
604            let s = match returning.as_ref() {
605                Some(returning) => {
606                    let (row_members, row_construct) = construct_row(returning);
607                    quote! { {
608                        use ::sqlx::Arguments as _;
609                        let _ = std::include_bytes!(#sp);
610                        #(#errors; )*
611                        #args_tokens
612
613                        struct Row {
614                            #(#row_members),*
615                        };
616                        sqlx::__query_with_result(#q, query_args).map(|row|
617                            Row{
618                                #(#row_construct),*
619                            }
620                        )
621                    }}
622                }
623                None => quote! { {
624                    use ::sqlx::Arguments as _;
625                    const _SCHEMA_HASH: u64 = #schema_hash;
626                    #(#errors; )*
627                    #args_tokens
628                    sqlx::__query_with_result(#q, query_args)
629                }
630                },
631            };
632            s.into()
633        }
634        qusql_type::StatementType::Insert {
635            arguments,
636            returning,
637            ..
638        } => {
639            let (args_tokens, q) = quote_args(
640                &mut errors,
641                &query.query,
642                query.last_span,
643                &query.args,
644                arguments,
645                dialect,
646            );
647            let s = match returning.as_ref() {
648                Some(returning) => {
649                    let (row_members, row_construct) = construct_row(returning);
650                    quote! { {
651                        use ::sqlx::Arguments as _;
652                        let _ = std::include_bytes!(#sp);
653                        #(#errors; )*
654                        #args_tokens
655
656                        struct Row {
657                            #(#row_members),*
658                        };
659                        sqlx::__query_with_result(#q, query_args).map(|row|
660                            Row{
661                                #(#row_construct),*
662                            }
663                        )
664                    }}
665                }
666                None => quote! { {
667                    use ::sqlx::Arguments as _;
668                    const _SCHEMA_HASH: u64 = #schema_hash;
669                    #(#errors; )*
670                    #args_tokens
671                    sqlx::__query_with_result(#q, query_args)
672                }
673                },
674            };
675            s.into()
676        }
677        qusql_type::StatementType::Update {
678            arguments,
679            returning,
680        } => {
681            let (args_tokens, q) = quote_args(
682                &mut errors,
683                &query.query,
684                query.last_span,
685                &query.args,
686                arguments,
687                dialect,
688            );
689
690            let s = match returning.as_ref() {
691                Some(returning) => {
692                    let (row_members, row_construct) = construct_row(returning);
693                    quote! { {
694                        use ::sqlx::Arguments as _;
695                        let _ = std::include_bytes!(#sp);
696                        #(#errors; )*
697                        #args_tokens
698
699                        struct Row {
700                            #(#row_members),*
701                        };
702                        sqlx::__query_with_result(#q, query_args).map(|row|
703                            Row{
704                                #(#row_construct),*
705                            }
706                        )
707                    }}
708                }
709                None => quote! { {
710                    use ::sqlx::Arguments as _;
711                    const _SCHEMA_HASH: u64 = #schema_hash;
712                    #(#errors; )*
713                    #args_tokens
714                    sqlx::__query_with_result(#q, query_args)
715                }
716                },
717            };
718            s.into()
719        }
720        qusql_type::StatementType::Replace {
721            arguments,
722            returning,
723        } => {
724            let (args_tokens, q) = quote_args(
725                &mut errors,
726                &query.query,
727                query.last_span,
728                &query.args,
729                arguments,
730                dialect,
731            );
732            let s = match returning.as_ref() {
733                Some(returning) => {
734                    let (row_members, row_construct) = construct_row(returning);
735                    quote! { {
736                        use ::sqlx::Arguments as _;
737                        const _SCHEMA_HASH: u64 = #schema_hash;
738                        let _ = std::include_bytes!(#sp);
739                        #(#errors; )*
740                        #args_tokens
741
742                        struct Row {
743                            #(#row_members),*
744                        };
745                        sqlx::__query_with_result(#q, query_args).map(|row|
746                            Row{
747                                #(#row_construct),*
748                            }
749                        )
750                    }}
751                }
752                None => quote! { {
753                    use ::sqlx::Arguments as _;
754                    #(#errors; )*
755                    #args_tokens
756                    sqlx::__query_with_result(#q, query_args)
757                }
758                },
759            };
760            s.into()
761        }
762        qusql_type::StatementType::Truncate => {
763            errors.push(
764                syn::Error::new(query.query_span, "TRUNCATE not supported in query!")
765                    .to_compile_error(),
766            );
767            quote! { {
768                #(#errors; )*
769                todo!("truncate")
770            }}
771            .into()
772        }
773        qusql_type::StatementType::Call { .. } => {
774            errors.push(
775                syn::Error::new(query.query_span, "CALL not supported in query!")
776                    .to_compile_error(),
777            );
778            quote! { {
779                #(#errors; )*
780                todo!("call")
781            }}
782            .into()
783        }
784        qusql_type::StatementType::Transaction => {
785            errors.push(
786                syn::Error::new(
787                    query.query_span,
788                    "Transaction control not supported in query!",
789                )
790                .to_compile_error(),
791            );
792            quote! { {
793                #(#errors; )*
794                todo!("transaction")
795            }}
796            .into()
797        }
798        qusql_type::StatementType::Set => {
799            errors.push(
800                syn::Error::new(query.query_span, "SET not supported in query!").to_compile_error(),
801            );
802            quote! { {
803                #(#errors; )*
804                todo!("set")
805            }}
806            .into()
807        }
808        qusql_type::StatementType::Lock => {
809            errors.push(
810                syn::Error::new(query.query_span, "LOCK not supported in query!")
811                    .to_compile_error(),
812            );
813            quote! { {
814                #(#errors; )*
815                todo!("lock")
816            }}
817            .into()
818        }
819        qusql_type::StatementType::Invalid => {
820            let s = quote! { {
821                #(#errors; )*;
822                todo!("Invalid")
823            }};
824            s.into()
825        }
826    }
827}
828
829/// Fill in row values in a query_as struct
830fn construct_row2(columns: &[SelectTypeColumn]) -> Vec<proc_macro2::TokenStream> {
831    let mut row_construct = Vec::new();
832    for (i, c) in columns.iter().enumerate() {
833        let mut t = match c.type_.t {
834            qusql_type::Type::U8 => quote! {u8},
835            qusql_type::Type::I8 => quote! {i8},
836            qusql_type::Type::U16 => quote! {u16},
837            qusql_type::Type::I16 => quote! {i16},
838            qusql_type::Type::U24 => quote! {u32},
839            qusql_type::Type::I24 => quote! {i32},
840            qusql_type::Type::U32 => quote! {u32},
841            qusql_type::Type::I32 => quote! {i32},
842            qusql_type::Type::U64 => quote! {u64},
843            qusql_type::Type::I64 => quote! {i64},
844            qusql_type::Type::Base(qusql_type::BaseType::Any) => todo!("from_any"),
845            qusql_type::Type::Base(qusql_type::BaseType::Bool) => quote! {bool},
846            qusql_type::Type::Base(qusql_type::BaseType::Bytes) => quote! {Vec<u8>},
847            qusql_type::Type::Base(qusql_type::BaseType::Date) => quote! {chrono::NaiveDate},
848            qusql_type::Type::Base(qusql_type::BaseType::DateTime) => {
849                quote! {chrono::NaiveDateTime}
850            }
851            qusql_type::Type::Base(qusql_type::BaseType::Float) => quote! {f64},
852            qusql_type::Type::Base(qusql_type::BaseType::Integer) => quote! {i64},
853            qusql_type::Type::Base(qusql_type::BaseType::String) => quote! {String},
854            qusql_type::Type::Base(qusql_type::BaseType::Time) => todo!("from_time"),
855            qusql_type::Type::Base(qusql_type::BaseType::TimeInterval) => {
856                todo!("from_time_interval")
857            }
858            qusql_type::Type::Base(qusql_type::BaseType::TimeStamp) => {
859                quote! {sqlx::types::chrono::DateTime<sqlx::types::chrono::Utc>}
860            }
861            qusql_type::Type::Base(qusql_type::BaseType::Uuid) => {
862                quote! {qusql_sqlx_type::UuidValue}
863            }
864            qusql_type::Type::Null => todo!("from_null"),
865            qusql_type::Type::Invalid => quote! {i64},
866            qusql_type::Type::Enum(_) => quote! {String},
867            qusql_type::Type::Set(_) => quote! {String},
868            qusql_type::Type::Args(_, _) => todo!("from_args"),
869            qusql_type::Type::F32 => quote! {f32},
870            qusql_type::Type::F64 => quote! {f64},
871            qusql_type::Type::JSON => quote! {String},
872            qusql_type::Type::Geometry => quote! {Vec<u8>},
873            qusql_type::Type::Range(_) => quote! {Vec<u8>},
874            qusql_type::Type::Array(_) => quote! {qusql_sqlx_type::Any},
875        };
876        let name = match &c.name {
877            Some(v) => v,
878            None => continue,
879        };
880
881        // Handle sqlx "column!" convention: strip trailing ! and force not_null
882        let name_str = name.value;
883        let (name_str, force_not_null) = if let Some(stripped) = name_str.strip_suffix('!') {
884            (stripped, true)
885        } else {
886            (name_str, false)
887        };
888
889        let ident = String::from("r#") + name_str;
890        let ident: Ident = if let Ok(ident) = syn::parse_str(&ident) {
891            ident
892        } else {
893            // TODO error
894            //errors.push(syn::Error::new(span, String::from_utf8(out).unwrap()).to_compile_error().into());
895            continue;
896        };
897
898        if !c.type_.not_null && !force_not_null {
899            t = quote! {Option<#t>};
900        }
901        row_construct.push(quote! {
902            #ident: qusql_sqlx_type::arg_out::<#t, _, #i>(sqlx::Row::get(&row, #i))
903        });
904    }
905    row_construct
906}
907
908/// Parse result of a query_as! macro
909struct QueryAs {
910    /// Name of output type
911    as_: Ident,
912    /// The query to execute
913    query: String,
914    /// The span of the query to execute
915    query_span: Span,
916    /// The arguments to supply
917    args: Vec<Expr>,
918    /// The last span parsed
919    last_span: Span,
920}
921
922impl Parse for QueryAs {
923    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
924        let as_ = input.parse::<Ident>()?;
925        let _ = input.parse::<syn::token::Comma>()?;
926
927        let query_ = Punctuated::<LitStr, Token![+]>::parse_separated_nonempty(input)?;
928        let query: String = query_.iter().map(LitStr::value).collect();
929        let query_span = query_.span();
930
931        let mut last_span = query_span;
932        let mut args = Vec::new();
933        while !input.is_empty() {
934            let _ = input.parse::<syn::token::Comma>()?;
935            if input.is_empty() {
936                break;
937            }
938            let arg = input.parse::<Expr>()?;
939            last_span = arg.span();
940            args.push(arg);
941        }
942        Ok(Self {
943            as_,
944            query,
945            query_span,
946            args,
947            last_span,
948        })
949    }
950}
951
952/// A variant of query! which takes a path to an explicitly defined struct as the output type.
953///
954/// This lets you return the struct from a function or add your own trait implementations.
955#[proc_macro]
956pub fn query_as(input: TokenStream) -> TokenStream {
957    let query_as = syn::parse_macro_input!(input as QueryAs);
958    let cache = get_schemas();
959    let (schemas, dialect) = (cache.schemas.get(), &cache.dialect);
960    let options = TypeOptions::new()
961        .dialect(dialect.clone())
962        .arguments(match &dialect {
963            SQLDialect::MariaDB => SQLArguments::QuestionMark,
964            SQLDialect::Sqlite => SQLArguments::QuestionMark,
965            SQLDialect::PostgreSQL | SQLDialect::PostGIS => SQLArguments::Dollar,
966        })
967        .list_hack(true);
968    let mut issues = qusql_type::Issues::new(&query_as.query);
969    let stmt = type_statement(&schemas.0, &query_as.query, &mut issues, &options);
970
971    let mut errors = issues_to_errors(issues.into_vec(), &query_as.query, query_as.query_span);
972    match &stmt {
973        qusql_type::StatementType::Select { columns, arguments } => {
974            let (args_tokens, q) = quote_args(
975                &mut errors,
976                &query_as.query,
977                query_as.last_span,
978                &query_as.args,
979                arguments,
980                dialect,
981            );
982
983            let row_construct = construct_row2(columns);
984            let row = query_as.as_;
985            let s = quote! { {
986                use ::sqlx::Arguments as _;
987                #(#errors; )*
988                #args_tokens
989                sqlx::__query_with_result(#q, query_args).map(|row|
990                    #row{
991                        #(#row_construct),*
992                    }
993                )
994            }};
995            //println!("TOKENS: {}", s);
996            s.into()
997        }
998        qusql_type::StatementType::Delete { .. } => {
999            errors.push(
1000                syn::Error::new(query_as.query_span, "DELETE not support in query_as")
1001                    .to_compile_error(),
1002            );
1003            quote! { {
1004                #(#errors; )*
1005                todo!("delete")
1006            }}
1007            .into()
1008        }
1009        qusql_type::StatementType::Insert {
1010            returning: None, ..
1011        } => {
1012            errors.push(
1013                syn::Error::new(
1014                    query_as.query_span,
1015                    "INSERT without RETURNING not support in query_as",
1016                )
1017                .to_compile_error(),
1018            );
1019            quote! { {
1020                #(#errors; )*
1021                todo!("insert")
1022            }}
1023            .into()
1024        }
1025        qusql_type::StatementType::Insert {
1026            arguments,
1027            returning: Some(returning),
1028            ..
1029        } => {
1030            let (args_tokens, q) = quote_args(
1031                &mut errors,
1032                &query_as.query,
1033                query_as.last_span,
1034                &query_as.args,
1035                arguments,
1036                dialect,
1037            );
1038
1039            let row_construct = construct_row2(returning);
1040            let row = query_as.as_;
1041            let s = quote! { {
1042                use ::sqlx::Arguments as _;
1043                #(#errors; )*
1044                #args_tokens
1045                sqlx::__query_with_result(#q, query_args).map(|row|
1046                    #row{
1047                        #(#row_construct),*
1048                    }
1049                )
1050            }};
1051            s.into()
1052        }
1053        qusql_type::StatementType::Update {
1054            returning: None, ..
1055        } => {
1056            errors.push(
1057                syn::Error::new(
1058                    query_as.query_span,
1059                    "UPDATE without RETURNING not support in query_as",
1060                )
1061                .to_compile_error(),
1062            );
1063            quote! { {
1064                #(#errors; )*
1065                todo!("update")
1066            }}
1067            .into()
1068        }
1069        qusql_type::StatementType::Update {
1070            arguments,
1071            returning: Some(returning),
1072            ..
1073        } => {
1074            let (args_tokens, q) = quote_args(
1075                &mut errors,
1076                &query_as.query,
1077                query_as.last_span,
1078                &query_as.args,
1079                arguments,
1080                dialect,
1081            );
1082
1083            let row_construct = construct_row2(returning);
1084            let row = query_as.as_;
1085            let s = quote! { {
1086                use ::sqlx::Arguments as _;
1087                #(#errors; )*
1088                #args_tokens
1089                sqlx::__query_with_result(#q, query_args).map(|row|
1090                    #row{
1091                        #(#row_construct),*
1092                    }
1093                )
1094            }};
1095            s.into()
1096        }
1097        qusql_type::StatementType::Replace {
1098            returning: None, ..
1099        } => {
1100            errors.push(
1101                syn::Error::new(
1102                    query_as.query_span,
1103                    "REPLACE without RETURNING not support in query_as",
1104                )
1105                .to_compile_error(),
1106            );
1107            quote! { {
1108                #(#errors; )*
1109                todo!("replace")
1110            }}
1111            .into()
1112        }
1113        qusql_type::StatementType::Replace {
1114            arguments,
1115            returning: Some(returning),
1116            ..
1117        } => {
1118            let (args_tokens, q) = quote_args(
1119                &mut errors,
1120                &query_as.query,
1121                query_as.last_span,
1122                &query_as.args,
1123                arguments,
1124                dialect,
1125            );
1126
1127            let row_construct = construct_row2(returning);
1128            let row = query_as.as_;
1129            let s = quote! { {
1130                use ::sqlx::Arguments as _;
1131                #(#errors; )*
1132                #args_tokens
1133                sqlx::__query_with_result(#q, query_args).map(|row|
1134                    #row{
1135                        #(#row_construct),*
1136                    }
1137                )
1138            }};
1139            s.into()
1140        }
1141        qusql_type::StatementType::Truncate => {
1142            errors.push(
1143                syn::Error::new(query_as.query_span, "TRUNCATE not supported in query_as!")
1144                    .to_compile_error(),
1145            );
1146            quote! { {
1147                #(#errors; )*
1148                todo!("truncate")
1149            }}
1150            .into()
1151        }
1152        qusql_type::StatementType::Call { .. } => {
1153            errors.push(
1154                syn::Error::new(query_as.query_span, "CALL not supported in query_as!")
1155                    .to_compile_error(),
1156            );
1157            quote! { {
1158                #(#errors; )*
1159                todo!("call")
1160            }}
1161            .into()
1162        }
1163        qusql_type::StatementType::Transaction => {
1164            errors.push(
1165                syn::Error::new(
1166                    query_as.query_span,
1167                    "Transaction control not supported in query_as!",
1168                )
1169                .to_compile_error(),
1170            );
1171            quote! { {
1172                #(#errors; )*
1173                todo!("transaction")
1174            }}
1175            .into()
1176        }
1177        qusql_type::StatementType::Set => {
1178            errors.push(
1179                syn::Error::new(query_as.query_span, "SET not supported in query_as!")
1180                    .to_compile_error(),
1181            );
1182            quote! { {
1183                #(#errors; )*
1184                todo!("set")
1185            }}
1186            .into()
1187        }
1188        qusql_type::StatementType::Lock => {
1189            errors.push(
1190                syn::Error::new(query_as.query_span, "LOCK not supported in query_as!")
1191                    .to_compile_error(),
1192            );
1193            quote! { {
1194                #(#errors; )*
1195                todo!("lock")
1196            }}
1197            .into()
1198        }
1199        qusql_type::StatementType::Invalid => quote! { {
1200            #(#errors; )*;
1201            todo!("invalid")
1202        }}
1203        .into(),
1204    }
1205}