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        // Pre-wrapped to satisfy `cargo fmt --check` at the default 100-col
38        // max_width — the single-line form exceeds the limit and rustfmt
39        // rewrites it.
40        let header = format!(
41            "// Code generated by sqlc-gen-sqlx v{}. DO NOT EDIT.\n\
42             // sqlc version: {}\n\n\
43             #![allow(\n    dead_code,\n    reason = \"generated queries may expose items a caller does not use\"\n)]\n\n",
44            self.plugin_version, self.sqlc_version,
45        );
46
47        Ok(format!("{header}{formatted}"))
48    }
49}
50
51#[cfg(test)]
52mod tests {
53    use super::*;
54
55    #[test]
56    fn emits_valid_rust() {
57        let mut e = FileEmitter::new("0.0.0-test", "sqlc-test");
58        e.push(quote::quote! { pub struct Foo { pub x: i32 } });
59        let code = e.finish().unwrap();
60        assert!(code.contains("pub struct Foo"));
61        assert!(code.contains("pub x: i32"));
62        // Header comment present
63        assert!(code.contains("DO NOT EDIT"));
64        assert!(code.contains("sqlc-gen-sqlx vsqlc-test"));
65        assert!(code.contains("sqlc version: 0.0.0-test"));
66        assert!(code.contains(
67            "#![allow(\n    dead_code,\n    reason = \"generated queries may expose items a caller does not use\"\n)]"
68        ));
69    }
70
71    #[test]
72    fn emits_empty_file_with_header() {
73        let e = FileEmitter::new("1.0.0", "sqlc");
74        let code = e.finish().unwrap();
75        assert!(code.contains("DO NOT EDIT"));
76    }
77}