Skip to main content

sqlc_gen_sqlx/
emit.rs

1use crate::error::Error;
2
3/// Collects `TokenStream` fragments and formats them into a Rust source file
4/// using `prettyplease`.
5pub struct FileEmitter {
6    sqlc_version: String,
7    plugin_version: String,
8    items: Vec<proc_macro2::TokenStream>,
9}
10
11impl FileEmitter {
12    pub fn new(sqlc_version: &str, plugin_version: &str) -> Self {
13        Self {
14            sqlc_version: sqlc_version.to_string(),
15            plugin_version: plugin_version.to_string(),
16            items: Vec::new(),
17        }
18    }
19
20    /// Append a token stream fragment (a struct, impl block, fn, const, etc.).
21    pub fn push(&mut self, tokens: proc_macro2::TokenStream) {
22        self.items.push(tokens);
23    }
24
25    /// Render all items into a formatted Rust source string.
26    pub fn finish(self) -> Result<String, Error> {
27        use quote::quote;
28
29        let items = &self.items;
30        let combined = quote! { #(#items)* };
31
32        let file: syn::File = syn::parse2(combined)
33            .map_err(|e| Error::Codegen(format!("token stream parse failed: {e}")))?;
34
35        let formatted = prettyplease::unparse(&file);
36
37        let header = format!(
38            "// Code generated by sqlc-gen-sqlx v{}. DO NOT EDIT.\n\
39             // sqlc version: {}\n\n\
40             #![allow(dead_code, reason = \"generated queries may expose items a caller does not use\")]\n\n",
41            self.plugin_version, self.sqlc_version,
42        );
43
44        Ok(format!("{header}{formatted}"))
45    }
46}
47
48#[cfg(test)]
49mod tests {
50    use super::*;
51
52    #[test]
53    fn emits_valid_rust() {
54        let mut e = FileEmitter::new("0.0.0-test", "sqlc-test");
55        e.push(quote::quote! { pub struct Foo { pub x: i32 } });
56        let code = e.finish().unwrap();
57        assert!(code.contains("pub struct Foo"));
58        assert!(code.contains("pub x: i32"));
59        // Header comment present
60        assert!(code.contains("DO NOT EDIT"));
61        assert!(code.contains("sqlc-gen-sqlx vsqlc-test"));
62        assert!(code.contains("sqlc version: 0.0.0-test"));
63        assert!(code.contains(
64            "#![allow(dead_code, reason = \"generated queries may expose items a caller does not use\")]"
65        ));
66    }
67
68    #[test]
69    fn emits_empty_file_with_header() {
70        let e = FileEmitter::new("1.0.0", "sqlc");
71        let code = e.finish().unwrap();
72        assert!(code.contains("DO NOT EDIT"));
73    }
74}