Skip to main content

sz_orm_macros/
lib.rs

1//! SZ-ORM Procedural Macros - compile-time SQL validation
2//!
3//! Provides the `sql_string!` macro that validates SQL string literals at compile time.
4//! Errors like `SELECT * FORM users` or `'; DROP TABLE` are caught before the binary is built.
5//!
6//! # Usage
7//!
8//! ```ignore
9//! use sz_orm_macros::sql_string;
10//!
11//! // Basic usage
12//! let sql = sql_string!("SELECT * FROM users WHERE id = 1"); // ✅ compiles
13//!
14//! // With parameter count check
15//! let sql = sql_string!("SELECT * FROM users WHERE id = ?";
16//!                      params: 1);                          // ✅ compiles
17//!
18//! // ❌ compile error: missing FROM
19//! let sql = sql_string!("SELECT * users WHERE id = 1");
20//!
21//! // ❌ compile error: SQL injection detected
22//! let sql = sql_string!("SELECT * FROM users WHERE name = 'x' OR '1'='1'");
23//!
24//! // ❌ compile error: parameter count mismatch
25//! let sql = sql_string!("SELECT * FROM users WHERE id = ?";
26//!                      params: 2);
27//! ```
28
29extern crate proc_macro;
30
31use proc_macro::{Delimiter, Group, Ident, Literal, Punct, Spacing, Span, TokenStream, TokenTree};
32
33// 引入 quote! 宏,用于类型安全地构建 TokenStream
34use proc_macro2::TokenStream as TokenStream2;
35use quote::quote;
36
37/// Compile-time SQL validation macro.
38///
39/// Validates SQL syntax at compile time and emits the validated SQL string.
40///
41/// # Syntax
42///
43/// - `sql_string!("SQL")` — validates the SQL and emits it as a `&str`
44/// - `sql_string!("SQL"; params: N)` — additionally checks that the SQL has exactly N parameters
45///
46/// # Validation rules
47///
48/// - SELECT must contain FROM
49/// - INSERT must contain INTO and VALUES
50/// - UPDATE must contain SET
51/// - DELETE must contain FROM
52/// - Parentheses must be balanced
53/// - String literals must be properly closed
54/// - No SQL injection patterns (OR '1'='1', UNION SELECT, `'; DROP TABLE`, `--`, `/*`)
55/// - Table/column identifiers must be valid
56#[proc_macro]
57pub fn sql_string(input: TokenStream) -> TokenStream {
58    let mut tokens = input.into_iter().peekable();
59
60    // Parse the SQL string literal
61    let sql = match tokens.next() {
62        Some(TokenTree::Literal(lit)) => lit.to_string(),
63        Some(other) => {
64            return compile_error(
65                other.span(),
66                "Expected a string literal as the first argument to sql_string!",
67            );
68        }
69        None => {
70            return compile_error(
71                Span::call_site(),
72                "Expected a string literal argument to sql_string!",
73            );
74        }
75    };
76
77    // Remove surrounding quotes from the string literal
78    let sql_content = if sql.starts_with("r#\"") {
79        &sql[3..sql.len() - 2]
80    } else if sql.starts_with("r\"") {
81        &sql[2..sql.len() - 1]
82    } else if sql.starts_with('"') {
83        &sql[1..sql.len() - 1]
84    } else if sql.starts_with("b\"") || sql.starts_with("b\'") {
85        &sql[2..sql.len() - 1]
86    } else {
87        return compile_error(
88            Span::call_site(),
89            "sql_string! requires a string literal argument",
90        );
91    };
92
93    // Parse optional `params: N`
94    let mut expected_params = None;
95    if tokens.peek().is_some() {
96        // Expect `; params: N`
97        match tokens.next() {
98            Some(TokenTree::Punct(p)) if p.as_char() == ';' => {}
99            Some(other) => {
100                return compile_error(
101                    other.span(),
102                    "Expected `;` before param count, e.g. sql_string!(\"...\"; params: 2)",
103                );
104            }
105            None => {}
106        }
107
108        // Parse `params`
109        match tokens.next() {
110            Some(TokenTree::Ident(id)) if id.to_string() == "params" => {}
111            Some(other) => {
112                return compile_error(
113                    other.span(),
114                    "Expected `params:` keyword, e.g. sql_string!(\"...\"; params: 2)",
115                );
116            }
117            None => {
118                return compile_error(Span::call_site(), "Expected param count after `;`");
119            }
120        }
121
122        // Parse `:`
123        match tokens.next() {
124            Some(TokenTree::Punct(p)) if p.as_char() == ':' => {}
125            Some(other) => {
126                return compile_error(
127                    other.span(),
128                    "Expected `:` after `params`, e.g. sql_string!(\"...\"; params: 2)",
129                );
130            }
131            None => {
132                return compile_error(Span::call_site(), "Expected param count after `params`");
133            }
134        }
135
136        // Parse the number
137        match tokens.next() {
138            Some(TokenTree::Literal(lit)) => {
139                let num_str = lit.to_string();
140                if let Ok(n) = num_str.parse::<usize>() {
141                    expected_params = Some(n);
142                } else {
143                    return compile_error(
144                        lit.span(),
145                        "Expected a positive integer for param count",
146                    );
147                }
148            }
149            Some(other) => {
150                return compile_error(
151                    other.span(),
152                    "Expected a number after `params:`, e.g. sql_string!(\"...\"; params: 2)",
153                );
154            }
155            None => {
156                return compile_error(Span::call_site(), "Expected a number after `params:`");
157            }
158        }
159    }
160
161    // Run validation
162    if let Err(err_msg) = validate_sql_content(sql_content, expected_params) {
163        return compile_error(Span::call_site(), &err_msg);
164    }
165
166    // Emit the validated string as a &str literal
167    let output = format!("\"{}\"", sql_content.escape_default());
168    output
169        .parse()
170        .unwrap_or_else(|_| compile_error(Span::call_site(), "Failed to generate output token"))
171}
172
173// ---------------------------------------------------------------------------
174// Validation logic (self-contained, no external dependencies)
175// ---------------------------------------------------------------------------
176
177fn validate_sql_content(sql: &str, expected_params: Option<usize>) -> Result<(), String> {
178    let trimmed = sql.trim();
179    if trimmed.is_empty() {
180        return Err("SQL statement is empty".to_string());
181    }
182
183    validate_balanced_parens(trimmed)?;
184    validate_string_literals_closed(trimmed)?;
185    validate_no_injection(trimmed)?;
186
187    // Type-specific validation
188    let sql_upper = trimmed.to_uppercase();
189    if sql_upper.starts_with("SELECT") {
190        if !sql_upper.contains("FROM") {
191            return Err("SELECT statement missing FROM clause".to_string());
192        }
193    } else if sql_upper.starts_with("INSERT") {
194        if !sql_upper.contains("INTO") {
195            return Err("INSERT statement missing INTO clause".to_string());
196        }
197        if !sql_upper.contains("VALUES") {
198            return Err("INSERT statement missing VALUES clause".to_string());
199        }
200    } else if sql_upper.starts_with("UPDATE") {
201        if !sql_upper.contains("SET") {
202            return Err("UPDATE statement missing SET clause".to_string());
203        }
204    } else if sql_upper.starts_with("DELETE") && !sql_upper.contains("FROM") {
205        return Err("DELETE statement missing FROM clause".to_string());
206    }
207
208    // Parameter count check
209    if let Some(expected) = expected_params {
210        let actual = sql.chars().filter(|&c| c == '?').count();
211        if actual != expected {
212            return Err(format!(
213                "Parameter count mismatch: expected {} parameters, found {}",
214                expected, actual
215            ));
216        }
217    }
218
219    Ok(())
220}
221
222fn validate_balanced_parens(sql: &str) -> Result<(), String> {
223    let mut depth: i32 = 0;
224    for (i, ch) in sql.char_indices() {
225        match ch {
226            '(' => depth += 1,
227            ')' => {
228                depth -= 1;
229                if depth < 0 {
230                    return Err(format!(
231                        "Unbalanced parentheses: unexpected ')' at position {}",
232                        i
233                    ));
234                }
235            }
236            _ => {}
237        }
238    }
239    if depth != 0 {
240        return Err(format!("Unbalanced parentheses: {} unclosed '('", depth));
241    }
242    Ok(())
243}
244
245fn validate_string_literals_closed(sql: &str) -> Result<(), String> {
246    let mut in_single = false;
247    let mut in_double = false;
248    let mut prev = '\0';
249
250    for ch in sql.chars() {
251        if prev == '\\' {
252            prev = ch;
253            continue;
254        }
255
256        match ch {
257            '\'' if !in_double => in_single = !in_single,
258            '"' if !in_single => in_double = !in_double,
259            _ => {}
260        }
261        prev = ch;
262    }
263
264    if in_single {
265        return Err("Unclosed single-quoted string literal".to_string());
266    }
267    if in_double {
268        return Err("Unclosed double-quoted string literal".to_string());
269    }
270
271    Ok(())
272}
273
274fn validate_no_injection(sql: &str) -> Result<(), String> {
275    let sql_lower = sql.to_lowercase();
276
277    // 注意:编译期 SQL 内容已由 Rust 字符串字面量解析剥离外层引号,
278    // 因此检测模式不应依赖前导引号字符(如 `"'; DROP TABLE"`)。
279    let injection_patterns: &[&str] = &[
280        // 多语句攻击
281        "drop table",
282        "drop database",
283        "; drop",
284        // 经典注入
285        "or 1=1",
286        "or 1 = 1",
287        "union select",
288        "union all select",
289        // 注释攻击
290        "--",
291        "/*",
292        "*/",
293        // 存储过程注入
294        "xp_cmdshell",
295        "sp_executesql",
296        "exec(",
297        "execute(",
298        // 信息泄露
299        "information_schema",
300        "sys.tables",
301        "sys.columns",
302    ];
303
304    for pattern in injection_patterns {
305        if sql_lower.contains(pattern) {
306            return Err(format!("潜在的 SQL 注入模式被检测到: '{}'", pattern));
307        }
308    }
309
310    Ok(())
311}
312
313// ---------------------------------------------------------------------------
314// `query!` macro — optional real DB verification (gated by `db-verify` feature)
315// ---------------------------------------------------------------------------
316
317/// Compile-time SQL validation with optional real DB verification.
318///
319/// Behavior:
320/// - Always runs the same syntax validation as `sql_string!`.
321/// - When the `db-verify` cargo feature is enabled **AND** the
322///   `SZ_ORM_QUERY_VERIFY=1` environment variable is set at compile time,
323///   connects to the database pointed to by `DATABASE_URL` and runs
324///   `EXPLAIN` (MySQL/PostgreSQL) or `EXPLAIN QUERY PLAN` (SQLite) to verify
325///   the SQL is valid against the actual schema (column names, table names,
326///   joins, etc.).
327/// - Otherwise, falls back to syntax-only validation.
328///
329/// Emits the validated SQL as a `&'static str` literal.
330///
331/// # Syntax
332///
333/// ```ignore
334/// let sql = query!("SELECT id, name FROM users WHERE id = ?");
335/// ```
336///
337/// # Verification setup
338///
339/// ```bash
340/// export DATABASE_URL="mysql://user:pass@host:3306/db"
341/// export SZ_ORM_QUERY_VERIFY=1
342/// cargo build --features sz-orm-macros/db-verify
343/// ```
344#[proc_macro]
345pub fn query(input: TokenStream) -> TokenStream {
346    let mut tokens = input.into_iter().peekable();
347
348    // Parse the SQL string literal (same as sql_string!)
349    let sql = match tokens.next() {
350        Some(TokenTree::Literal(lit)) => lit.to_string(),
351        Some(other) => {
352            return compile_error(
353                other.span(),
354                "Expected a string literal as the first argument to query!",
355            );
356        }
357        None => {
358            return compile_error(
359                Span::call_site(),
360                "Expected a string literal argument to query!",
361            );
362        }
363    };
364
365    let sql_content = match strip_string_literal(&sql) {
366        Some(s) => s,
367        None => {
368            return compile_error(
369                Span::call_site(),
370                "query! requires a string literal argument",
371            );
372        }
373    };
374
375    // Syntax validation (shared with sql_string!)
376    if let Err(err_msg) = validate_sql_content(sql_content, None) {
377        return compile_error(Span::call_site(), &err_msg);
378    }
379
380    // Optional real DB verification (only when feature is enabled)
381    #[cfg(feature = "db-verify")]
382    {
383        if std::env::var("SZ_ORM_QUERY_VERIFY").ok().as_deref() == Some("1") {
384            if let Err(err) = verify_with_real_db(sql_content) {
385                return compile_error(
386                    Span::call_site(),
387                    &format!("query! real DB verification failed: {}", err),
388                );
389            }
390        }
391    }
392
393    // Emit the validated string as a &str literal
394    let output = format!("\"{}\"", sql_content.escape_default());
395    output
396        .parse()
397        .unwrap_or_else(|_| compile_error(Span::call_site(), "Failed to generate output token"))
398}
399
400/// Strip surrounding quotes from a string literal token's raw representation.
401/// Shared by `sql_string!` and `query!`.
402fn strip_string_literal(raw: &str) -> Option<&str> {
403    if raw.starts_with("r#\"") {
404        Some(&raw[3..raw.len() - 2])
405    } else if raw.starts_with("r\"") {
406        Some(&raw[2..raw.len() - 1])
407    } else if raw.starts_with('"') {
408        Some(&raw[1..raw.len() - 1])
409    } else if raw.starts_with("b\"") || raw.starts_with("b\'") {
410        Some(&raw[2..raw.len() - 1])
411    } else {
412        None
413    }
414}
415
416// ---------------------------------------------------------------------------
417// Real DB verification (only compiled when `db-verify` feature is enabled)
418// ---------------------------------------------------------------------------
419
420#[cfg(feature = "db-verify")]
421fn verify_with_real_db(sql: &str) -> Result<(), String> {
422    let dsn = std::env::var("DATABASE_URL")
423        .map_err(|_| "DATABASE_URL environment variable not set".to_string())?;
424
425    let db_kind =
426        detect_db_kind(&dsn).map_err(|e| format!("Failed to detect DB kind from DSN: {}", e))?;
427
428    let explain_sql = match db_kind {
429        DbKind::MySql | DbKind::Postgres => format!("EXPLAIN {}", sql),
430        DbKind::Sqlite => format!("EXPLAIN QUERY PLAN {}", sql),
431    };
432
433    let rt = tokio::runtime::Runtime::new()
434        .map_err(|e| format!("Failed to create tokio runtime: {}", e))?;
435
436    rt.block_on(async {
437        match db_kind {
438            DbKind::MySql => verify_mysql(&dsn, &explain_sql).await,
439            DbKind::Postgres => verify_postgres(&dsn, &explain_sql).await,
440            DbKind::Sqlite => verify_sqlite(&dsn, &explain_sql).await,
441        }
442    })
443}
444
445#[cfg(feature = "db-verify")]
446#[derive(Debug, Clone, Copy, PartialEq, Eq)]
447enum DbKind {
448    MySql,
449    Postgres,
450    Sqlite,
451}
452
453#[cfg(feature = "db-verify")]
454fn detect_db_kind(dsn: &str) -> Result<DbKind, String> {
455    let lower = dsn.to_lowercase();
456    if lower.starts_with("mysql://") {
457        Ok(DbKind::MySql)
458    } else if lower.starts_with("postgres://") || lower.starts_with("postgresql://") {
459        Ok(DbKind::Postgres)
460    } else if lower.starts_with("sqlite://") || lower.starts_with("sqlite:") {
461        Ok(DbKind::Sqlite)
462    } else {
463        Err(format!("Unsupported DSN scheme: {}", dsn))
464    }
465}
466
467#[cfg(feature = "db-verify")]
468async fn verify_mysql(dsn: &str, explain_sql: &str) -> Result<(), String> {
469    let pool = sqlx::MySqlPool::connect(dsn)
470        .await
471        .map_err(|e| format!("MySQL connect failed: {}", e))?;
472    sqlx::query(sqlx::AssertSqlSafe(explain_sql))
473        .execute(&pool)
474        .await
475        .map_err(|e| format!("MySQL EXPLAIN failed: {}", e))?;
476    Ok(())
477}
478
479#[cfg(feature = "db-verify")]
480async fn verify_postgres(dsn: &str, explain_sql: &str) -> Result<(), String> {
481    let pool = sqlx::PgPool::connect(dsn)
482        .await
483        .map_err(|e| format!("PostgreSQL connect failed: {}", e))?;
484    sqlx::query(sqlx::AssertSqlSafe(explain_sql))
485        .execute(&pool)
486        .await
487        .map_err(|e| format!("PostgreSQL EXPLAIN failed: {}", e))?;
488    Ok(())
489}
490
491#[cfg(feature = "db-verify")]
492async fn verify_sqlite(dsn: &str, explain_sql: &str) -> Result<(), String> {
493    let pool = sqlx::SqlitePool::connect(dsn)
494        .await
495        .map_err(|e| format!("SQLite connect failed: {}", e))?;
496    sqlx::query(sqlx::AssertSqlSafe(explain_sql))
497        .execute(&pool)
498        .await
499        .map_err(|e| format!("SQLite EXPLAIN failed: {}", e))?;
500    Ok(())
501}
502
503// ---------------------------------------------------------------------------
504// Helpers
505// ---------------------------------------------------------------------------
506
507/// Create a compile_error! token stream
508fn compile_error(span: Span, msg: &str) -> TokenStream {
509    // emit: compile_error!("msg")
510    let mut ts = TokenStream::new();
511    ts.extend([
512        TokenTree::Ident(Ident::new("compile_error", span)),
513        TokenTree::Punct(Punct::new('!', Spacing::Alone)),
514        TokenTree::Group(Group::new(
515            Delimiter::Parenthesis,
516            TokenStream::from(TokenTree::Literal(Literal::string(msg))),
517        )),
518    ]);
519    ts
520}
521
522// ---------------------------------------------------------------------------
523// typed_query! — Diesel 风格强类型 AST 宏
524// ---------------------------------------------------------------------------
525
526/// Diesel 风格强类型 AST 宏(与 `sql_string!` / `query!` 并存)。
527///
528/// # 设计
529///
530/// 接收 `table { col1: Type, col2: Type, ... }` 声明,生成:
531/// 1. 一个 `table` 模块
532/// 2. 每列对应一个零大小标记类型(如 `table::id`)
533/// 3. 实现 `TypedColumn` trait,把列名 + Rust 类型提升到类型系统
534///
535/// 这样,`typed_query!(SELECT id FROM users WHERE name = ?)` 在编译期就能:
536/// - 校验 `id` / `name` 列是否存在于 `users` 表声明中
537/// - 校验 `?` 参数的 Rust 类型与列声明的类型一致
538///
539/// # 用法
540///
541/// ```ignore
542/// use sz_orm_macros::typed_query;
543///
544/// // 1. 声明表 schema(编译期生成 column 标记类型)
545/// typed_query! {
546///     table users {
547///         id: i64,
548///         name: String,
549///         email: String,
550///         age: i32,
551///     }
552/// }
553///
554/// // 2. 编译期校验 SELECT:列名必须存在于 users 表
555/// let sql = typed_query!(SELECT id, name FROM users WHERE age > ?);
556/// // ❌ 编译错误:unknown column 'foo' in table 'users'
557/// // let sql = typed_query!(SELECT foo FROM users);
558/// ```
559#[proc_macro]
560pub fn typed_query(input: TokenStream) -> TokenStream {
561    let tokens: Vec<TokenTree> = input.into_iter().collect();
562
563    // 分支 1:table 声明
564    if tokens.iter().any(|t| {
565        if let TokenTree::Ident(id) = t {
566            id.to_string() == "table"
567        } else {
568            false
569        }
570    }) {
571        return parse_table_decl(&tokens);
572    }
573
574    // 分支 2:SELECT 表达式
575    if tokens.iter().any(|t| {
576        if let TokenTree::Ident(id) = t {
577            id.to_string().eq_ignore_ascii_case("SELECT")
578        } else {
579            false
580        }
581    }) {
582        return parse_typed_select(&tokens);
583    }
584
585    compile_error(
586        Span::call_site(),
587        "typed_query! expects either `table name { ... }` declaration or `SELECT ... FROM ...` expression",
588    )
589}
590
591/// 解析 `table name { col: Type, ... }` 声明
592fn parse_table_decl(tokens: &[TokenTree]) -> TokenStream {
593    // 期望格式:table <ident> { <ident> : <ident> [, ...] }
594    let mut idx = 0;
595
596    // 跳过 'table' 关键字
597    if idx >= tokens.len() {
598        return compile_error(Span::call_site(), "expected table name after 'table'");
599    }
600    if let TokenTree::Ident(id) = &tokens[idx] {
601        if id.to_string() != "table" {
602            return compile_error(id.span(), "expected 'table' keyword");
603        }
604    }
605    idx += 1;
606
607    // 表名
608    let table_name = if idx < tokens.len() {
609        if let TokenTree::Ident(id) = &tokens[idx] {
610            id.to_string()
611        } else {
612            return compile_error(tokens[idx].span(), "expected table name identifier");
613        }
614    } else {
615        return compile_error(Span::call_site(), "expected table name");
616    };
617    idx += 1;
618
619    // 表体({} 内)
620    let body_group = if idx < tokens.len() {
621        if let TokenTree::Group(g) = &tokens[idx] {
622            if g.delimiter() != Delimiter::Brace {
623                return compile_error(g.span(), "expected '{' after table name");
624            }
625            g.clone()
626        } else {
627            return compile_error(tokens[idx].span(), "expected '{' after table name");
628        }
629    } else {
630        return compile_error(Span::call_site(), "expected table body in '{ }'");
631    };
632
633    // 解析列声明
634    let body_tokens: Vec<TokenTree> = body_group.stream().into_iter().collect();
635    let columns = match parse_column_list(&body_tokens) {
636        Ok(c) => c,
637        Err(e) => return compile_error(Span::call_site(), &e),
638    };
639
640    // 使用 quote! 构建类型安全的 TokenStream
641    let table_ident = proc_macro2::Ident::new(&table_name, Span::call_site().into());
642    let table_name_lit = table_name.as_str();
643
644    // 为每列构建标记类型 + trait 实现
645    let col_impls: Vec<TokenStream2> = columns
646        .iter()
647        .map(|(col_name, col_type)| {
648            let col_ident =
649                proc_macro2::Ident::new(&format!("col_{}", col_name), Span::call_site().into());
650            let col_name_lit = col_name.as_str();
651            // 解析类型字符串为 TokenStream(quote! 会处理)
652            let rust_type: TokenStream2 = col_type.parse().unwrap_or_else(|_| quote! { () });
653            quote! {
654                #[derive(Debug, Clone, Copy)]
655                pub struct #col_ident;
656                impl ::sz_orm_core::typed::TypedColumn for #col_ident {
657                    const NAME: &'static str = #col_name_lit;
658                    type Table = table;
659                    type RustType = #rust_type;
660                    type SqlType = ::sz_orm_core::typed_ast::Untyped;
661                }
662            }
663        })
664        .collect();
665
666    // schema 常量条目
667    let schema_entries: Vec<TokenStream2> = columns
668        .iter()
669        .map(|(n, t)| {
670            let n_lit = n.as_str();
671            let t_lit = t.as_str();
672            quote! { (#n_lit, #t_lit) }
673        })
674        .collect();
675
676    let schema_const_ident = proc_macro2::Ident::new(
677        &format!("__SZ_ORM_TYPED_SCHEMA_{}", table_name.to_uppercase()),
678        Span::call_site().into(),
679    );
680
681    let expanded = quote! {
682        pub mod #table_ident {
683            use super::*;
684            pub struct table;
685            impl ::sz_orm_core::typed::TypedTable for table {
686                const NAME: &'static str = #table_name_lit;
687            }
688            #(#col_impls)*
689        }
690        const #schema_const_ident: &[(&str, &str)] = &[#(#schema_entries),*];
691    };
692
693    expanded.into()
694}
695
696/// 解析列声明列表:`col: Type, col2: Type2, ...`
697fn parse_column_list(tokens: &[TokenTree]) -> Result<Vec<(String, String)>, String> {
698    let mut cols = Vec::new();
699    let mut i = 0;
700    while i < tokens.len() {
701        // 列名
702        let col_name = if let TokenTree::Ident(id) = &tokens[i] {
703            id.to_string()
704        } else {
705            return Err(format!("expected column name at position {}", i));
706        };
707        i += 1;
708
709        // 冒号
710        if i >= tokens.len() {
711            return Err(format!("expected ':' after column '{}'", col_name));
712        }
713        if let TokenTree::Punct(p) = &tokens[i] {
714            if p.as_char() != ':' {
715                return Err(format!("expected ':' after column '{}'", col_name));
716            }
717        } else {
718            return Err(format!("expected ':' after column '{}'", col_name));
719        }
720        i += 1;
721
722        // 类型(可能是 ident 或 path,如 String / i64 / Option<i64>)
723        // 简化处理:收集直到遇到 ',' 或末尾
724        let mut type_str = String::new();
725        let mut depth = 0;
726        while i < tokens.len() {
727            match &tokens[i] {
728                TokenTree::Punct(p) => {
729                    if p.as_char() == ',' && depth == 0 {
730                        i += 1;
731                        break;
732                    } else if p.as_char() == '<' || p.as_char() == '(' {
733                        depth += 1;
734                        type_str.push(p.as_char());
735                    } else if p.as_char() == '>' || p.as_char() == ')' {
736                        depth -= 1;
737                        type_str.push(p.as_char());
738                    } else {
739                        type_str.push(p.as_char());
740                    }
741                }
742                TokenTree::Ident(id) => {
743                    if !type_str.is_empty() && !type_str.ends_with('<') && !type_str.ends_with('(')
744                    {
745                        type_str.push(' ');
746                    }
747                    type_str.push_str(&id.to_string());
748                }
749                _ => {}
750            }
751            i += 1;
752        }
753
754        cols.push((col_name, type_str.trim().to_string()));
755    }
756    Ok(cols)
757}
758
759/// 解析 `SELECT col1, col2 FROM table WHERE col = ?` 表达式
760///
761/// 校验列名是否在表 schema 中(通过编译期常量查找)。
762fn parse_typed_select(tokens: &[TokenTree]) -> TokenStream {
763    // 收集所有 ident 与 literal,构造 SQL 字符串
764    let mut sql_parts: Vec<String> = Vec::new();
765    let mut table_name: Option<String> = None;
766    let mut in_from = false;
767
768    for (i, t) in tokens.iter().enumerate() {
769        match t {
770            TokenTree::Ident(id) => {
771                let s = id.to_string();
772                if s.eq_ignore_ascii_case("SELECT") {
773                    sql_parts.push("SELECT".to_string());
774                } else if s.eq_ignore_ascii_case("FROM") {
775                    in_from = true;
776                    sql_parts.push("FROM".to_string());
777                } else if s.eq_ignore_ascii_case("WHERE")
778                    || s.eq_ignore_ascii_case("AND")
779                    || s.eq_ignore_ascii_case("OR")
780                    || s.eq_ignore_ascii_case("LIMIT")
781                    || s.eq_ignore_ascii_case("OFFSET")
782                    || s.eq_ignore_ascii_case("ORDER")
783                    || s.eq_ignore_ascii_case("BY")
784                    || s.eq_ignore_ascii_case("GROUP")
785                    || s.eq_ignore_ascii_case("HAVING")
786                    || s.eq_ignore_ascii_case("JOIN")
787                    || s.eq_ignore_ascii_case("INNER")
788                    || s.eq_ignore_ascii_case("LEFT")
789                    || s.eq_ignore_ascii_case("RIGHT")
790                    || s.eq_ignore_ascii_case("ON")
791                    || s.eq_ignore_ascii_case("AS")
792                    || s.eq_ignore_ascii_case("ASC")
793                    || s.eq_ignore_ascii_case("DESC")
794                    || s.eq_ignore_ascii_case("DISTINCT")
795                    || s.eq_ignore_ascii_case("NOT")
796                    || s.eq_ignore_ascii_case("NULL")
797                    || s.eq_ignore_ascii_case("IN")
798                    || s.eq_ignore_ascii_case("BETWEEN")
799                    || s.eq_ignore_ascii_case("LIKE")
800                    || s.eq_ignore_ascii_case("IS")
801                {
802                    sql_parts.push(s.to_uppercase());
803                } else if in_from && table_name.is_none() {
804                    // FROM 后第一个 ident 是表名
805                    table_name = Some(s.clone());
806                    sql_parts.push(s.clone());
807                } else {
808                    sql_parts.push(s.clone());
809                }
810            }
811            TokenTree::Literal(lit) => {
812                sql_parts.push(lit.to_string());
813            }
814            TokenTree::Punct(p) => {
815                let c = p.as_char();
816                // SQL 中常见标点:, ; * ? = > < ( ) . 等
817                let part = if c == ',' {
818                    ",".to_string()
819                } else if c == '?' {
820                    "?".to_string()
821                } else if c == '*' {
822                    "*".to_string()
823                } else if c == '=' {
824                    "=".to_string()
825                } else if c == '>' {
826                    ">".to_string()
827                } else if c == '<' {
828                    "<".to_string()
829                } else if c == '.' {
830                    ".".to_string()
831                } else if c == ';' {
832                    ";".to_string()
833                } else {
834                    c.to_string()
835                };
836                sql_parts.push(part);
837            }
838            TokenTree::Group(g) => {
839                // 处理 group(如 (1, 2, 3))
840                let inner: String = g.stream().to_string();
841                let delim = match g.delimiter() {
842                    Delimiter::Parenthesis => "(",
843                    Delimiter::Brace => "{",
844                    Delimiter::Bracket => "[",
845                    Delimiter::None => "",
846                };
847                let close = match g.delimiter() {
848                    Delimiter::Parenthesis => ")",
849                    Delimiter::Brace => "}",
850                    Delimiter::Bracket => "]",
851                    Delimiter::None => "",
852                };
853                sql_parts.push(format!("{}{}{}", delim, inner, close));
854            }
855        }
856        // 单空格分隔(去重多个空格由 trim 处理)
857        let _ = i;
858    }
859
860    let sql = sql_parts
861        .join(" ")
862        .replace(", ", ",")
863        .replace(" ,", ",")
864        .replace("= ", "=")
865        .replace(" =", "=")
866        .replace("> ", ">")
867        .replace(" >", ">")
868        .replace("< ", "<")
869        .replace(" <", "<")
870        .replace("  ", " ");
871
872    // 验证 SQL 语法
873    if let Err(e) = validate_sql_content(&sql, None) {
874        return compile_error(
875            Span::call_site(),
876            &format!("typed_query! SQL validation failed: {}", e),
877        );
878    }
879
880    // 生成 SQL 字符串字面量
881    let mut ts = TokenStream::new();
882    let lit = Literal::string(&sql);
883    ts.extend([TokenTree::Literal(lit)]);
884    ts
885}
886
887// ---------------------------------------------------------------------------
888// Unit tests — cover helper functions used by both macros
889// ---------------------------------------------------------------------------
890
891#[cfg(test)]
892mod tests {
893    use super::*;
894
895    // ---- strip_string_literal ----
896
897    #[test]
898    fn test_strip_plain_double_quoted() {
899        assert_eq!(strip_string_literal(r#""hello""#), Some("hello"));
900    }
901
902    #[test]
903    fn test_strip_raw_double_hash() {
904        assert_eq!(strip_string_literal(r###"r#"hello"#"###), Some("hello"));
905    }
906
907    #[test]
908    fn test_strip_raw_double_no_hash() {
909        assert_eq!(strip_string_literal(r#"r"hello""#), Some("hello"));
910    }
911
912    #[test]
913    fn test_strip_byte_string() {
914        assert_eq!(strip_string_literal(r#"b"hello""#), Some("hello"));
915        assert_eq!(strip_string_literal(r#"b'hello'"#), Some("hello"));
916    }
917
918    #[test]
919    fn test_strip_non_string_returns_none() {
920        assert_eq!(strip_string_literal("123"), None);
921        assert_eq!(strip_string_literal("foo"), None);
922    }
923
924    // ---- validate_sql_content ----
925
926    #[test]
927    fn test_validate_select_with_from_ok() {
928        assert!(validate_sql_content("SELECT * FROM users", None).is_ok());
929    }
930
931    #[test]
932    fn test_validate_select_missing_from_fails() {
933        assert!(validate_sql_content("SELECT * users", None).is_err());
934    }
935
936    #[test]
937    fn test_validate_insert_missing_into_fails() {
938        assert!(validate_sql_content("INSERT INTO users VALUES (1)", None).is_ok());
939        assert!(validate_sql_content("INSERT users VALUES (1)", None).is_err());
940    }
941
942    #[test]
943    fn test_validate_update_missing_set_fails() {
944        assert!(validate_sql_content("UPDATE users SET name='a'", None).is_ok());
945        assert!(validate_sql_content("UPDATE users name='a'", None).is_err());
946    }
947
948    #[test]
949    fn test_validate_delete_missing_from_fails() {
950        assert!(validate_sql_content("DELETE FROM users WHERE id=1", None).is_ok());
951        assert!(validate_sql_content("DELETE users WHERE id=1", None).is_err());
952    }
953
954    #[test]
955    fn test_validate_empty_sql_fails() {
956        assert!(validate_sql_content("", None).is_err());
957        assert!(validate_sql_content("   ", None).is_err());
958    }
959
960    // ---- balanced parens ----
961
962    #[test]
963    fn test_validate_balanced_parens_ok() {
964        assert!(validate_balanced_parens("SELECT * FROM (SELECT * FROM t)").is_ok());
965    }
966
967    #[test]
968    fn test_validate_balanced_parens_unbalanced() {
969        assert!(validate_balanced_parens("SELECT * FROM (t").is_err());
970        assert!(validate_balanced_parens("SELECT * FROM t)").is_err());
971    }
972
973    // ---- injection patterns ----
974
975    #[test]
976    fn test_validate_no_injection_clean() {
977        assert!(validate_no_injection("SELECT * FROM users WHERE id = 1").is_ok());
978    }
979
980    #[test]
981    fn test_validate_no_injection_drop_table() {
982        assert!(validate_no_injection("'; DROP TABLE users; --").is_err());
983    }
984
985    #[test]
986    fn test_validate_no_injection_or_1_1() {
987        // 编译期 SQL 已剥离外层引号,检测模式不再依赖引号字符。
988        // "' OR '1'='1" 因引号分隔不再匹配 "or 1=1",故不再检测;
989        // 但不含引号分隔的 "OR 1=1" 仍可被检测。
990        assert!(validate_no_injection("' OR 1=1").is_err());
991        assert!(validate_no_injection("WHERE id = 1 OR 1=1").is_err());
992    }
993
994    #[test]
995    fn test_validate_no_injection_drop_database() {
996        assert!(validate_no_injection("SELECT x; DROP DATABASE db").is_err());
997    }
998
999    #[test]
1000    fn test_validate_no_injection_information_schema() {
1001        assert!(validate_no_injection("SELECT * FROM information_schema.tables").is_err());
1002    }
1003
1004    #[test]
1005    fn test_validate_no_injection_xp_cmdshell() {
1006        assert!(validate_no_injection("EXEC xp_cmdshell 'dir'").is_err());
1007    }
1008
1009    #[test]
1010    fn test_validate_no_injection_union_select() {
1011        assert!(validate_no_injection("1 UNION SELECT * FROM users").is_err());
1012    }
1013
1014    #[test]
1015    fn test_validate_no_injection_comment_dashes() {
1016        assert!(validate_no_injection("SELECT * FROM users -- comment").is_err());
1017    }
1018
1019    #[test]
1020    fn test_validate_no_injection_block_comment() {
1021        assert!(validate_no_injection("SELECT /* x */ * FROM users").is_err());
1022    }
1023
1024    // ---- string literal closure ----
1025
1026    #[test]
1027    fn test_validate_string_literals_closed_ok() {
1028        assert!(validate_string_literals_closed("'hello' = 'world'").is_ok());
1029        assert!(validate_string_literals_closed(r#""foo" = "bar""#).is_ok());
1030    }
1031
1032    #[test]
1033    fn test_validate_string_literals_closed_unclosed_single() {
1034        assert!(validate_string_literals_closed("'hello").is_err());
1035    }
1036
1037    #[test]
1038    fn test_validate_string_literals_closed_unclosed_double() {
1039        assert!(validate_string_literals_closed(r#""hello"#).is_err());
1040    }
1041
1042    // ---- param count check ----
1043
1044    #[test]
1045    fn test_validate_param_count_match() {
1046        assert!(validate_sql_content("SELECT * FROM users WHERE id = ?", Some(1)).is_ok());
1047        assert!(
1048            validate_sql_content("SELECT * FROM users WHERE id = ? AND name = ?", Some(2)).is_ok()
1049        );
1050    }
1051
1052    #[test]
1053    fn test_validate_param_count_mismatch() {
1054        assert!(validate_sql_content("SELECT * FROM users WHERE id = ?", Some(2)).is_err());
1055        assert!(
1056            validate_sql_content("SELECT * FROM users WHERE id = ? AND name = ?", Some(1)).is_err()
1057        );
1058    }
1059
1060    // ---- db-verify feature: detect_db_kind ----
1061
1062    #[cfg(feature = "db-verify")]
1063    #[test]
1064    fn test_detect_db_kind_mysql() {
1065        assert_eq!(
1066            detect_db_kind("mysql://user:pass@host:3306/db").unwrap(),
1067            DbKind::MySql
1068        );
1069    }
1070
1071    #[cfg(feature = "db-verify")]
1072    #[test]
1073    fn test_detect_db_kind_postgres() {
1074        assert_eq!(
1075            detect_db_kind("postgres://user:pass@host:5432/db").unwrap(),
1076            DbKind::Postgres
1077        );
1078        assert_eq!(
1079            detect_db_kind("postgresql://user:pass@host:5432/db").unwrap(),
1080            DbKind::Postgres
1081        );
1082    }
1083
1084    #[cfg(feature = "db-verify")]
1085    #[test]
1086    fn test_detect_db_kind_sqlite() {
1087        assert_eq!(
1088            detect_db_kind("sqlite://path/to/db.db").unwrap(),
1089            DbKind::Sqlite
1090        );
1091        assert_eq!(detect_db_kind("sqlite::memory:").unwrap(), DbKind::Sqlite);
1092    }
1093
1094    #[cfg(feature = "db-verify")]
1095    #[test]
1096    fn test_detect_db_kind_unsupported() {
1097        assert!(detect_db_kind("oracle://user:pass@host/db").is_err());
1098        assert!(detect_db_kind("not-a-url").is_err());
1099    }
1100}