tenant.openagents/omega

1use proc_macro::{Delimiter, Span, TokenStream, TokenTree};
2use syn::Error;
3
4#[cfg(not(any(target_os = "linux", target_os = "freebsd")))]
5static SQLITE: std::sync::LazyLock<sqlez::thread_safe_connection::ThreadSafeConnection> =
6    std::sync::LazyLock::new(|| {
7        sqlez::thread_safe_connection::ThreadSafeConnection::new(
8            ":memory:",
9            false,
10            None,
11            Some(sqlez::thread_safe_connection::locking_queue()),
12        )
13    });
14
15#[proc_macro]
16pub fn sql(tokens: TokenStream) -> TokenStream {
17    let (spans, sql) = make_sql(tokens);
18
19    #[cfg(not(any(target_os = "linux", target_os = "freebsd")))]
20    let error = SQLITE.sql_has_syntax_error(sql.trim());
21
22    #[cfg(any(target_os = "linux", target_os = "freebsd"))]
23    let error: Option<(String, usize)> = None;
24
25    let formatted_sql = sqlformat::format(&sql, &sqlformat::QueryParams::None, Default::default());
26
27    if let Some((error, error_offset)) = error {
28        create_error(spans, error_offset, error, &formatted_sql)
29    } else {
30        format!("r#\"{}\"#", &formatted_sql).parse().unwrap()
31    }
32}
33
34fn create_error(
35    spans: Vec<(usize, Span)>,
36    error_offset: usize,
37    error: String,
38    formatted_sql: &String,
39) -> TokenStream {
40    let error_span = spans
41        .into_iter()
42        .skip_while(|(offset, _)| offset <= &error_offset)
43        .map(|(_, span)| span)
44        .next()
45        .unwrap_or_else(Span::call_site);
46    let error_text = format!("Sql Error: {}\nFor Query: {}", error, formatted_sql);
47    TokenStream::from(Error::new(error_span.into(), error_text).into_compile_error())
48}
49
50fn make_sql(tokens: TokenStream) -> (Vec<(usize, Span)>, String) {
51    let mut sql_tokens = vec![];
52    flatten_stream(tokens, &mut sql_tokens);
53    // Lookup of spans by offset at the end of the token
54    let mut spans: Vec<(usize, Span)> = Vec::new();
55    let mut sql = String::new();
56    for (token_text, span) in sql_tokens {
57        sql.push_str(&token_text);
58        spans.push((sql.len(), span));
59    }
60    (spans, sql)
61}
62
63/// This method exists to normalize the representation of groups
64/// to always include spaces between tokens. This is why we don't use the usual .to_string().
65/// This allows our token search in token_at_offset to resolve
66/// ambiguity of '(tokens)' vs. '( token )', due to sqlite requiring byte offsets
67fn flatten_stream(tokens: TokenStream, result: &mut Vec<(String, Span)>) {
68    for token_tree in tokens.into_iter() {
69        match token_tree {
70            TokenTree::Group(group) => {
71                // push open delimiter
72                result.push((open_delimiter(group.delimiter()), group.span()));
73                // recurse
74                flatten_stream(group.stream(), result);
75                // push close delimiter
76                result.push((close_delimiter(group.delimiter()), group.span()));
77            }
78            TokenTree::Ident(ident) => {
79                result.push((format!("{} ", ident), ident.span()));
80            }
81            leaf_tree => result.push((leaf_tree.to_string(), leaf_tree.span())),
82        }
83    }
84}
85
86fn open_delimiter(delimiter: Delimiter) -> String {
87    match delimiter {
88        Delimiter::Parenthesis => "( ".to_string(),
89        Delimiter::Brace => "[ ".to_string(),
90        Delimiter::Bracket => "{ ".to_string(),
91        Delimiter::None => "".to_string(),
92    }
93}
94
95fn close_delimiter(delimiter: Delimiter) -> String {
96    match delimiter {
97        Delimiter::Parenthesis => " ) ".to_string(),
98        Delimiter::Brace => " ] ".to_string(),
99        Delimiter::Bracket => " } ".to_string(),
100        Delimiter::None => "".to_string(),
101    }
102}
103