Skip to main content

scythe_core/parser/
mod.rs

1use sqlparser::parser::Parser;
2
3use crate::dialect::SqlDialect;
4use crate::errors::ScytheError;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum QueryCommand {
8    One,
9    Opt,
10    Many,
11    Exec,
12    ExecResult,
13    ExecRows,
14    Batch,
15    Grouped,
16}
17
18impl std::fmt::Display for QueryCommand {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        match self {
21            QueryCommand::One => write!(f, "one"),
22            QueryCommand::Opt => write!(f, "opt"),
23            QueryCommand::Many => write!(f, "many"),
24            QueryCommand::Exec => write!(f, "exec"),
25            QueryCommand::ExecResult => write!(f, "exec_result"),
26            QueryCommand::ExecRows => write!(f, "exec_rows"),
27            QueryCommand::Batch => write!(f, "batch"),
28            QueryCommand::Grouped => write!(f, "grouped"),
29        }
30    }
31}
32
33impl QueryCommand {
34    fn from_str(s: &str) -> Result<Self, ScytheError> {
35        match s {
36            "one" => Ok(QueryCommand::One),
37            "opt" => Ok(QueryCommand::Opt),
38            "many" => Ok(QueryCommand::Many),
39            "exec" => Ok(QueryCommand::Exec),
40            "exec_result" => Ok(QueryCommand::ExecResult),
41            "exec_rows" => Ok(QueryCommand::ExecRows),
42            "batch" => Ok(QueryCommand::Batch),
43            "grouped" => Ok(QueryCommand::Grouped),
44            other => Err(ScytheError::invalid_annotation(format!(
45                "invalid @returns value: {other}"
46            ))),
47        }
48    }
49}
50
51#[derive(Debug, Clone, PartialEq, Eq)]
52pub struct ParamDoc {
53    pub name: String,
54    pub description: String,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq)]
58pub struct JsonMapping {
59    pub column: String,
60    pub rust_type: String,
61}
62
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub struct Annotations {
65    pub name: String,
66    pub command: QueryCommand,
67    pub param_docs: Vec<ParamDoc>,
68    pub nullable_overrides: Vec<String>,
69    pub nonnull_overrides: Vec<String>,
70    pub json_mappings: Vec<JsonMapping>,
71    pub deprecated: Option<String>,
72    pub optional_params: Vec<String>,
73    pub group_by: Option<String>,
74}
75
76#[derive(Debug)]
77pub struct Query {
78    pub name: String,
79    pub command: QueryCommand,
80    pub sql: String,
81    pub stmt: sqlparser::ast::Statement,
82    pub annotations: Annotations,
83}
84
85/// Parse a single annotated SQL query into a `Query` using the PostgreSQL dialect.
86pub fn parse_query(query_sql: &str) -> Result<Query, ScytheError> {
87    parse_query_with_dialect(query_sql, &SqlDialect::PostgreSQL)
88}
89
90/// Parse a single annotated SQL query into a `Query` using the specified dialect.
91pub fn parse_query_with_dialect(
92    query_sql: &str,
93    dialect: &SqlDialect,
94) -> Result<Query, ScytheError> {
95    let mut name: Option<String> = None;
96    let mut command: Option<QueryCommand> = None;
97    let mut param_docs = Vec::new();
98    let mut nullable_overrides = Vec::new();
99    let mut nonnull_overrides = Vec::new();
100    let mut json_mappings = Vec::new();
101    let mut deprecated: Option<String> = None;
102    let mut optional_params = Vec::new();
103    let mut group_by: Option<String> = None;
104
105    let mut sql_lines = Vec::new();
106
107    for line in query_sql.lines() {
108        let trimmed = line.trim();
109
110        // Check for annotation: "-- @..." or "--@..."
111        let annotation_body = if let Some(rest) = trimmed.strip_prefix("--") {
112            let rest = rest.trim_start();
113            rest.strip_prefix('@')
114        } else {
115            None
116        };
117
118        if let Some(body) = annotation_body {
119            // Parse the annotation keyword and value
120            let (keyword, value) = match body.find(|c: char| c.is_whitespace()) {
121                Some(pos) => (&body[..pos], body[pos..].trim()),
122                None => (body, ""),
123            };
124
125            match keyword.to_ascii_lowercase().as_str() {
126                "name" => {
127                    name = Some(value.to_string());
128                }
129                "returns" => {
130                    let cmd_str = value.strip_prefix(':').unwrap_or(value);
131                    command = Some(QueryCommand::from_str(cmd_str)?);
132                }
133                "param" => {
134                    // format: "<name>: <description>" or "<name>:<description>"
135                    if let Some(colon_pos) = value.find(':') {
136                        let param_name = value[..colon_pos].trim().to_string();
137                        let description = value[colon_pos + 1..].trim().to_string();
138                        param_docs.push(ParamDoc {
139                            name: param_name,
140                            description,
141                        });
142                    } else {
143                        param_docs.push(ParamDoc {
144                            name: value.to_string(),
145                            description: String::new(),
146                        });
147                    }
148                }
149                "nullable" => {
150                    for col in value.split(',') {
151                        let col = col.trim();
152                        if !col.is_empty() {
153                            nullable_overrides.push(col.to_string());
154                        }
155                    }
156                }
157                "nonnull" => {
158                    for col in value.split(',') {
159                        let col = col.trim();
160                        if !col.is_empty() {
161                            nonnull_overrides.push(col.to_string());
162                        }
163                    }
164                }
165                "json" => {
166                    // format: "<col> = <Type>"
167                    if let Some(eq_pos) = value.find('=') {
168                        let column = value[..eq_pos].trim().to_string();
169                        let rust_type = value[eq_pos + 1..].trim().to_string();
170                        json_mappings.push(JsonMapping { column, rust_type });
171                    }
172                }
173                "deprecated" => {
174                    deprecated = Some(value.to_string());
175                }
176                "group_by" => {
177                    group_by = Some(value.to_string());
178                }
179                "optional" => {
180                    for param in value.split(',') {
181                        let param = param.trim();
182                        if !param.is_empty() {
183                            optional_params.push(param.to_string());
184                        }
185                    }
186                }
187                _ => {
188                    // Unknown annotation — ignore or could error
189                }
190            }
191        } else {
192            sql_lines.push(line);
193        }
194    }
195
196    let name = name.ok_or_else(|| ScytheError::missing_annotation("name"))?;
197    let command = command.ok_or_else(|| ScytheError::missing_annotation("returns"))?;
198
199    if command == QueryCommand::Grouped && group_by.is_none() {
200        return Err(ScytheError::invalid_annotation(
201            "@returns :grouped requires a @group_by annotation (e.g. @group_by users.id)",
202        ));
203    }
204
205    let sql = sql_lines.join("\n").trim().to_string();
206
207    if sql.is_empty() {
208        return Err(ScytheError::syntax("empty SQL body"));
209    }
210
211    // Preprocess dialect-specific syntax before parsing:
212    //   * Oracle: strip `RETURNING ... INTO` output binds, convert `:N` → `?`.
213    //   * MSSQL: convert `OUTPUT INSERTED.*` → `RETURNING` for parsing,
214    //     convert `@pN` → `?` for parsing; keep original SQL for codegen.
215    //   * PostgreSQL: strip `WHERE …` between `ON CONFLICT (cols)` and `DO …`
216    //     for parsing (sqlparser-rs <= 0.61 doesn't recognise the
217    //     partial-index inference form); keep original SQL for codegen.
218    let (sql, parse_sql) = if *dialect == SqlDialect::Oracle {
219        let processed = preprocess_oracle_sql(&sql);
220        (processed.clone(), processed)
221    } else if *dialect == SqlDialect::MsSql {
222        // For codegen: only convert @pN → ? placeholders (keep OUTPUT syntax)
223        let codegen_sql = convert_mssql_placeholders(&sql);
224        // For parsing: also convert OUTPUT INSERTED → RETURNING
225        let parse_sql = preprocess_mssql_sql(&sql);
226        (codegen_sql, parse_sql)
227    } else if *dialect == SqlDialect::PostgreSQL {
228        let parse_sql = preprocess_postgres_sql(&sql);
229        (sql.clone(), parse_sql)
230    } else {
231        (sql.clone(), sql)
232    };
233
234    let parser_dialect = dialect.to_sqlparser_dialect();
235    let statements = Parser::parse_sql(parser_dialect.as_ref(), &parse_sql)
236        .map_err(|e| ScytheError::syntax(format!("syntax error: {}", e)))?;
237
238    if statements.len() != 1 {
239        // sqlparser may produce an extra empty statement from a trailing semicolon —
240        // filter those out by checking for exactly one non-empty statement.
241        let non_empty: Vec<_> = statements
242            .into_iter()
243            .filter(|s| {
244                !matches!(s, sqlparser::ast::Statement::Flush { .. }) && format!("{s}") != ""
245            })
246            .collect();
247        if non_empty.len() != 1 {
248            return Err(ScytheError::syntax("expected exactly one SQL statement"));
249        }
250        let stmt = non_empty
251            .into_iter()
252            .next()
253            .expect("filtered to exactly one statement");
254        let annotations = Annotations {
255            name: name.clone(),
256            command: command.clone(),
257            param_docs,
258            nullable_overrides,
259            nonnull_overrides,
260            json_mappings,
261            deprecated,
262            optional_params,
263            group_by: group_by.clone(),
264        };
265        return Ok(Query {
266            name,
267            command,
268            sql,
269            stmt,
270            annotations,
271        });
272    }
273
274    let stmt = statements
275        .into_iter()
276        .next()
277        .expect("filtered to exactly one statement");
278
279    let annotations = Annotations {
280        name: name.clone(),
281        command: command.clone(),
282        param_docs,
283        nullable_overrides,
284        nonnull_overrides,
285        json_mappings,
286        deprecated,
287        optional_params,
288        group_by,
289    };
290
291    Ok(Query {
292        name,
293        command,
294        sql,
295        stmt,
296        annotations,
297    })
298}
299
300/// Strip the `WHERE …` predicate that PostgreSQL allows between
301/// `ON CONFLICT (cols)` and `DO …` (the index-inference form for partial
302/// unique indexes). sqlparser-rs through 0.61 does not parse this construct;
303/// we lift it out for the parser and let the caller keep the original SQL
304/// for codegen + runtime, where Postgres validates it.
305fn preprocess_postgres_sql(sql: &str) -> String {
306    // Find a case-insensitive token sequence: ON CONFLICT ( ... ) WHERE ... DO
307    // We only touch the `WHERE ... DO` slice between the matching parenthesis
308    // close and the `DO` keyword, leaving the rest of the SQL untouched.
309    let upper = sql.to_uppercase();
310    let mut search_from = 0;
311    let mut result = String::with_capacity(sql.len());
312    let mut last = 0;
313    while let Some(rel) = upper[search_from..].find("ON CONFLICT") {
314        let on_conflict_pos = search_from + rel;
315        // Locate the column list `(...)` that follows ON CONFLICT.
316        let after_on_conflict = on_conflict_pos + "ON CONFLICT".len();
317        let bytes = sql.as_bytes();
318        let mut idx = after_on_conflict;
319        while idx < bytes.len() && bytes[idx].is_ascii_whitespace() {
320            idx += 1;
321        }
322        if idx >= bytes.len() || bytes[idx] != b'(' {
323            // ON CONFLICT ON CONSTRAINT … or bare ON CONFLICT — nothing to strip.
324            search_from = after_on_conflict;
325            continue;
326        }
327        // Find matching close paren.
328        let mut depth = 0i32;
329        let mut close = idx;
330        while close < bytes.len() {
331            match bytes[close] {
332                b'(' => depth += 1,
333                b')' => {
334                    depth -= 1;
335                    if depth == 0 {
336                        break;
337                    }
338                }
339                _ => {}
340            }
341            close += 1;
342        }
343        if depth != 0 {
344            // Unbalanced parens — bail and let sqlparser report the error.
345            return sql.to_string();
346        }
347        // Skip whitespace after the column list.
348        let mut after_cols = close + 1;
349        while after_cols < bytes.len() && bytes[after_cols].is_ascii_whitespace() {
350            after_cols += 1;
351        }
352        // Check for the WHERE keyword.
353        if upper[after_cols..].starts_with("WHERE") {
354            // Find the next DO keyword (must be word-bounded).
355            let where_start = after_cols;
356            let after_where = where_start + "WHERE".len();
357            if let Some(do_rel) = find_keyword(&upper[after_where..], "DO") {
358                let do_abs = after_where + do_rel;
359                // Emit everything up to (and including) the close paren + space,
360                // then skip the WHERE … and continue from `DO`.
361                result.push_str(&sql[last..after_cols]);
362                last = do_abs;
363                search_from = do_abs;
364                continue;
365            }
366        }
367        search_from = close + 1;
368    }
369    result.push_str(&sql[last..]);
370    result
371}
372
373/// Locate a whitespace-separated keyword in an uppercase haystack. Returns the
374/// byte offset of the keyword's start, or None if not found.
375fn find_keyword(haystack: &str, keyword: &str) -> Option<usize> {
376    let bytes = haystack.as_bytes();
377    let key = keyword.as_bytes();
378    let mut i = 0;
379    while i + key.len() <= bytes.len() {
380        if &bytes[i..i + key.len()] == key {
381            let prev_ok = i == 0 || !bytes[i - 1].is_ascii_alphanumeric();
382            let next = i + key.len();
383            let next_ok = next >= bytes.len() || !bytes[next].is_ascii_alphanumeric();
384            if prev_ok && next_ok {
385                return Some(i);
386            }
387        }
388        i += 1;
389    }
390    None
391}
392
393/// Preprocess Oracle SQL before parsing:
394/// 1. Strip `INTO :N, :N, ...` suffix from `RETURNING ... INTO` clauses.
395/// 2. Convert `:N` positional placeholders to `?` (universally supported).
396fn preprocess_oracle_sql(sql: &str) -> String {
397    // Strip Oracle RETURNING ... INTO clause (output bind variables)
398    // e.g. "INSERT ... RETURNING id, name INTO :4, :5" → "INSERT ... RETURNING id, name"
399    let sql = strip_returning_into(sql);
400
401    // Convert :N → ? (outside string literals)
402    let mut result = String::with_capacity(sql.len());
403    let mut chars = sql.chars().peekable();
404    while let Some(ch) = chars.next() {
405        if ch == '\'' {
406            // Skip string literals
407            result.push(ch);
408            while let Some(inner) = chars.next() {
409                result.push(inner);
410                if inner == '\'' {
411                    if chars.peek() == Some(&'\'') {
412                        result.push(chars.next().unwrap());
413                    } else {
414                        break;
415                    }
416                }
417            }
418        } else if ch == ':' && chars.peek().is_some_and(|c| c.is_ascii_digit()) {
419            // Convert :N → ?
420            result.push('?');
421            while chars.peek().is_some_and(|c| c.is_ascii_digit()) {
422                chars.next();
423            }
424        } else {
425            result.push(ch);
426        }
427    }
428    result
429}
430
431/// Convert MSSQL `@pN` positional placeholders to `?` (outside string literals).
432/// MsSqlDialect treats `@` as an identifier start, so `@p1` becomes an identifier
433/// rather than a `Placeholder` token — preprocessing normalises it to `?`.
434fn convert_mssql_placeholders(sql: &str) -> String {
435    let mut result = String::with_capacity(sql.len());
436    let mut chars = sql.chars().peekable();
437    while let Some(ch) = chars.next() {
438        if ch == '\'' {
439            // Skip string literals verbatim
440            result.push(ch);
441            while let Some(inner) = chars.next() {
442                result.push(inner);
443                if inner == '\'' {
444                    if chars.peek() == Some(&'\'') {
445                        // Escaped quote inside string literal
446                        result.push(chars.next().unwrap());
447                    } else {
448                        break;
449                    }
450                }
451            }
452        } else if ch == '@' && chars.peek().is_some_and(|c| *c == 'p' || *c == 'P') {
453            // Peek ahead: must be `@p` followed by at least one digit
454            let mut lookahead = chars.clone();
455            lookahead.next(); // consume the 'p'/'P'
456            if lookahead.peek().is_some_and(|c| c.is_ascii_digit()) {
457                // It is an `@pN` placeholder — consume `p` and all digits
458                chars.next(); // consume 'p'/'P'
459                while chars.peek().is_some_and(|c| c.is_ascii_digit()) {
460                    chars.next();
461                }
462                result.push('?');
463            } else {
464                result.push(ch);
465            }
466        } else {
467            result.push(ch);
468        }
469    }
470    result
471}
472
473/// Preprocess MSSQL SQL before parsing:
474/// 1. Strip `OUTPUT INSERTED.col, ...` clauses and convert to RETURNING
475/// 2. Convert `@pN` positional placeholders to `?`
476fn preprocess_mssql_sql(sql: &str) -> String {
477    // First pass: convert OUTPUT INSERTED.col to RETURNING col
478    let sql = strip_and_convert_mssql_output(sql);
479    // Second pass: convert @pN to ?
480    convert_mssql_placeholders(&sql)
481}
482
483/// Strip MSSQL `OUTPUT INSERTED.col1, INSERTED.col2, ...` from INSERT statements
484/// and convert it to a `RETURNING col1, col2, ...` clause.
485/// The OUTPUT clause appears between the column list and VALUES clause:
486///   INSERT INTO table (cols) OUTPUT INSERTED.col1, INSERTED.col2, ... VALUES (...)
487/// becomes:
488///   INSERT INTO table (cols) VALUES (...) RETURNING col1, col2, ...
489fn strip_and_convert_mssql_output(sql: &str) -> String {
490    // Case-insensitive search for OUTPUT keyword in INSERT statements
491    let upper = sql.to_uppercase();
492
493    // Only process INSERT statements with OUTPUT
494    if !upper.contains("INSERT") || !upper.contains("OUTPUT") {
495        return sql.to_string();
496    }
497
498    // Find the OUTPUT keyword
499    if let Some(output_pos) = find_word_position(&upper, "OUTPUT") {
500        // Check if this is actually part of an INSERT statement by finding INSERT before it
501        let before_output = &upper[..output_pos];
502        if !before_output.contains("INSERT") {
503            return sql.to_string();
504        }
505
506        // Look for the VALUES keyword after OUTPUT
507        let after_output = &upper[output_pos + "OUTPUT".len()..];
508        if let Some(values_offset) = find_word_position(after_output, "VALUES") {
509            let values_pos = output_pos + "OUTPUT".len() + values_offset;
510
511            // Extract the OUTPUT column list (between OUTPUT and VALUES)
512            let output_cols_str = &sql[output_pos + "OUTPUT".len()..values_pos];
513
514            // Parse column names: strip "INSERTED." prefix from each column name
515            let cols = parse_inserted_columns(output_cols_str);
516
517            if !cols.is_empty() {
518                // Build result: keep everything before OUTPUT, then VALUES clause,
519                // then RETURNING clause (before any trailing semicolon)
520                let before_output_sql = sql[..output_pos].trim_end();
521                let after_values = sql[values_pos..].trim_end();
522                let (values_body, trailing) = if let Some(stripped) = after_values.strip_suffix(';')
523                {
524                    (stripped, ";")
525                } else {
526                    (after_values, "")
527                };
528
529                return format!(
530                    "{}\n{} RETURNING {}{}",
531                    before_output_sql, values_body, cols, trailing
532                );
533            }
534        }
535    }
536
537    sql.to_string()
538}
539
540/// Find the position of a word (case-insensitive) in the text.
541/// The word must be a separate word, not part of another identifier.
542fn find_word_position(text: &str, word: &str) -> Option<usize> {
543    let mut pos = 0;
544    let word_len = word.len();
545    while let Some(idx) = text[pos..].find(word) {
546        let abs_idx = pos + idx;
547
548        // Check character before
549        let before_ok = abs_idx == 0
550            || !text
551                .as_bytes()
552                .get(abs_idx - 1)
553                .is_some_and(|&b| b.is_ascii_alphanumeric() || b == b'_');
554
555        // Check character after
556        let after_idx = abs_idx + word_len;
557        let after_ok = after_idx >= text.len()
558            || !text
559                .as_bytes()
560                .get(after_idx)
561                .is_some_and(|&b| b.is_ascii_alphanumeric() || b == b'_');
562
563        if before_ok && after_ok {
564            return Some(abs_idx);
565        }
566        pos = abs_idx + 1;
567    }
568    None
569}
570
571/// Parse INSERTED.col1, INSERTED.col2, ... and extract column names as "col1, col2, ..."
572fn parse_inserted_columns(output_str: &str) -> String {
573    let mut cols = Vec::new();
574
575    for part in output_str.split(',') {
576        let trimmed = part.trim();
577
578        // Try to extract column name after INSERTED.
579        if let Some(after_inserted) = trimmed
580            .strip_prefix("INSERTED.")
581            .or_else(|| trimmed.strip_prefix("inserted."))
582            .or_else(|| trimmed.strip_prefix("INSERTED"))
583            .or_else(|| trimmed.strip_prefix("inserted"))
584        {
585            let col_name = after_inserted.trim().to_string();
586            if !col_name.is_empty() {
587                cols.push(col_name);
588            }
589        }
590    }
591
592    cols.join(", ")
593}
594
595/// Strip the `INTO :N, :N, ...` suffix from an Oracle `RETURNING ... INTO` clause.
596fn strip_returning_into(sql: &str) -> String {
597    // Case-insensitive search for "INTO" after "RETURNING" at the end of the statement
598    let upper = sql.to_uppercase();
599    if let Some(ret_pos) = upper.rfind("RETURNING") {
600        let after_returning = &upper[ret_pos + "RETURNING".len()..];
601        if let Some(into_offset) = after_returning.find("INTO") {
602            let into_pos = ret_pos + "RETURNING".len() + into_offset;
603            // Keep everything before INTO, trim trailing whitespace/semicolons
604            let trimmed = sql[..into_pos].trim_end();
605            return trimmed.to_string();
606        }
607    }
608    sql.to_string()
609}
610
611#[cfg(test)]
612mod tests {
613    use super::*;
614    use crate::errors::ErrorCode;
615
616    fn parse(sql: &str) -> Result<Query, ScytheError> {
617        parse_query(sql)
618    }
619
620    #[test]
621    fn test_basic_parse() {
622        let input = "-- @name GetUsers\n-- @returns :many\nSELECT * FROM users;";
623        let q = parse(input).unwrap();
624        assert_eq!(q.name, "GetUsers");
625        assert_eq!(q.command, QueryCommand::Many);
626        assert!(q.sql.contains("SELECT"));
627    }
628
629    #[test]
630    fn test_all_command_types() {
631        let cases = vec![
632            (":one", QueryCommand::One),
633            (":many", QueryCommand::Many),
634            (":exec", QueryCommand::Exec),
635            (":exec_result", QueryCommand::ExecResult),
636            (":exec_rows", QueryCommand::ExecRows),
637        ];
638        for (tag, expected) in cases {
639            let input = format!("-- @name Q\n-- @returns {}\nSELECT 1", tag);
640            let q = parse(&input).unwrap();
641            assert_eq!(q.command, expected, "failed for {}", tag);
642        }
643    }
644
645    #[test]
646    fn test_case_insensitive_keywords() {
647        let input = "-- @Name GetUsers\n-- @RETURNS :many\nSELECT 1";
648        let q = parse(input).unwrap();
649        assert_eq!(q.name, "GetUsers");
650        assert_eq!(q.command, QueryCommand::Many);
651    }
652
653    #[test]
654    fn test_missing_name_errors() {
655        let input = "-- @returns :many\nSELECT 1";
656        let err = parse(input).unwrap_err();
657        assert_eq!(err.code, ErrorCode::MissingAnnotation);
658        assert!(err.message.contains("name"));
659    }
660
661    #[test]
662    fn test_missing_returns_errors() {
663        let input = "-- @name Foo\nSELECT 1";
664        let err = parse(input).unwrap_err();
665        assert_eq!(err.code, ErrorCode::MissingAnnotation);
666        assert!(err.message.contains("returns"));
667    }
668
669    #[test]
670    fn test_invalid_returns_value() {
671        let input = "-- @name Foo\n-- @returns :invalid\nSELECT 1";
672        let err = parse(input).unwrap_err();
673        assert_eq!(err.code, ErrorCode::InvalidAnnotation);
674    }
675
676    #[test]
677    fn test_empty_name_value() {
678        // An empty name is accepted by the parser (it stores "")
679        let input = "-- @name\n-- @returns :one\nSELECT 1";
680        let q = parse(input).unwrap();
681        assert_eq!(q.name, "");
682    }
683
684    #[test]
685    fn test_param_annotation() {
686        let input = "-- @name Foo\n-- @returns :one\n-- @param id: the user ID\nSELECT 1";
687        let q = parse(input).unwrap();
688        assert_eq!(q.annotations.param_docs.len(), 1);
689        assert_eq!(q.annotations.param_docs[0].name, "id");
690        assert_eq!(q.annotations.param_docs[0].description, "the user ID");
691    }
692
693    #[test]
694    fn test_param_no_description() {
695        let input = "-- @name Foo\n-- @returns :one\n-- @param id\nSELECT 1";
696        let q = parse(input).unwrap();
697        assert_eq!(q.annotations.param_docs.len(), 1);
698        assert_eq!(q.annotations.param_docs[0].name, "id");
699        assert_eq!(q.annotations.param_docs[0].description, "");
700    }
701
702    #[test]
703    fn test_nullable_annotation() {
704        let input = "-- @name Foo\n-- @returns :one\n-- @nullable col1, col2\nSELECT 1";
705        let q = parse(input).unwrap();
706        assert_eq!(q.annotations.nullable_overrides, vec!["col1", "col2"]);
707    }
708
709    #[test]
710    fn test_nonnull_annotation() {
711        let input = "-- @name Foo\n-- @returns :one\n-- @nonnull col1\nSELECT 1";
712        let q = parse(input).unwrap();
713        assert_eq!(q.annotations.nonnull_overrides, vec!["col1"]);
714    }
715
716    #[test]
717    fn test_json_annotation() {
718        let input = "-- @name Foo\n-- @returns :one\n-- @json data = EventData\nSELECT 1";
719        let q = parse(input).unwrap();
720        assert_eq!(q.annotations.json_mappings.len(), 1);
721        assert_eq!(q.annotations.json_mappings[0].column, "data");
722        assert_eq!(q.annotations.json_mappings[0].rust_type, "EventData");
723    }
724
725    #[test]
726    fn test_deprecated_annotation() {
727        let input = "-- @name Foo\n-- @returns :one\n-- @deprecated Use V2\nSELECT 1";
728        let q = parse(input).unwrap();
729        assert_eq!(q.annotations.deprecated, Some("Use V2".to_string()));
730    }
731
732    #[test]
733    fn test_sql_syntax_error() {
734        let input = "-- @name Foo\n-- @returns :one\nSELCT * FROM users";
735        let err = parse(input).unwrap_err();
736        assert_eq!(err.code, ErrorCode::SyntaxError);
737    }
738
739    #[test]
740    fn test_trailing_semicolon() {
741        let input = "-- @name Foo\n-- @returns :one\nSELECT 1;";
742        let q = parse(input).unwrap();
743        assert_eq!(q.name, "Foo");
744    }
745
746    #[test]
747    fn test_multiple_statements_error() {
748        let input = "-- @name Foo\n-- @returns :one\nSELECT 1; SELECT 2;";
749        let err = parse(input).unwrap_err();
750        assert_eq!(err.code, ErrorCode::SyntaxError);
751    }
752
753    #[test]
754    fn test_sql_preserved_without_annotations() {
755        let input = "-- @name Foo\n-- @returns :one\nSELECT id, name FROM users WHERE id = $1";
756        let q = parse(input).unwrap();
757        assert_eq!(q.sql, "SELECT id, name FROM users WHERE id = $1");
758    }
759
760    #[test]
761    fn test_returns_without_colon_prefix() {
762        let input = "-- @name Foo\n-- @returns many\nSELECT 1";
763        let q = parse(input).unwrap();
764        assert_eq!(q.command, QueryCommand::Many);
765    }
766
767    #[test]
768    fn test_batch_command() {
769        let input = "-- @name Foo\n-- @returns :batch\nSELECT 1";
770        let q = parse(input).unwrap();
771        assert_eq!(q.command, QueryCommand::Batch);
772    }
773
774    #[test]
775    fn test_grouped_command_with_group_by() {
776        let input = "-- @name GetUsersWithOrders\n-- @returns :grouped\n-- @group_by users.id\nSELECT u.id, u.name FROM users u JOIN orders o ON o.user_id = u.id";
777        let q = parse(input).unwrap();
778        assert_eq!(q.command, QueryCommand::Grouped);
779        assert_eq!(q.annotations.group_by, Some("users.id".to_string()));
780    }
781
782    #[test]
783    fn test_grouped_command_without_group_by_errors() {
784        let input = "-- @name Foo\n-- @returns :grouped\nSELECT 1";
785        let err = parse(input).unwrap_err();
786        assert_eq!(err.code, ErrorCode::InvalidAnnotation);
787        assert!(err.message.contains("@group_by"));
788    }
789
790    #[test]
791    fn test_group_by_without_grouped_is_ignored() {
792        let input = "-- @name Foo\n-- @returns :many\n-- @group_by users.id\nSELECT 1";
793        let q = parse(input).unwrap();
794        assert_eq!(q.command, QueryCommand::Many);
795        assert_eq!(q.annotations.group_by, Some("users.id".to_string()));
796    }
797
798    #[test]
799    fn test_preprocess_postgres_strips_partial_index_where() {
800        let sql = "INSERT INTO billing_events (project_id, stripe_event_id) \
801                   VALUES ($1, $2) \
802                   ON CONFLICT (stripe_event_id) WHERE stripe_event_id IS NOT NULL DO NOTHING";
803        let cleaned = preprocess_postgres_sql(sql);
804        assert!(
805            !cleaned
806                .to_uppercase()
807                .contains("WHERE STRIPE_EVENT_ID IS NOT NULL"),
808            "WHERE clause must be stripped between ON CONFLICT cols and DO; got: {cleaned}"
809        );
810        assert!(
811            cleaned
812                .to_uppercase()
813                .contains("ON CONFLICT (STRIPE_EVENT_ID) DO NOTHING")
814        );
815        // sqlparser must accept the cleaned form.
816        sqlparser::parser::Parser::parse_sql(&sqlparser::dialect::PostgreSqlDialect {}, &cleaned)
817            .expect("cleaned SQL should parse");
818    }
819
820    #[test]
821    fn test_preprocess_postgres_no_op_when_no_partial_clause() {
822        let sql = "INSERT INTO t (a) VALUES ($1) ON CONFLICT (a) DO UPDATE SET a = EXCLUDED.a";
823        assert_eq!(preprocess_postgres_sql(sql), sql);
824    }
825
826    #[test]
827    fn test_preprocess_postgres_leaves_on_conflict_on_constraint_alone() {
828        let sql = "INSERT INTO t (a) VALUES ($1) ON CONFLICT ON CONSTRAINT t_a_uidx DO NOTHING";
829        assert_eq!(preprocess_postgres_sql(sql), sql);
830    }
831
832    #[test]
833    fn test_preprocess_postgres_handles_compound_index_cols() {
834        let sql = "INSERT INTO t (a, b) VALUES ($1, $2) \
835                   ON CONFLICT (a, b) WHERE a IS NOT NULL AND b > 0 DO UPDATE SET b = EXCLUDED.b";
836        let cleaned = preprocess_postgres_sql(sql);
837        assert!(
838            cleaned
839                .to_uppercase()
840                .contains("ON CONFLICT (A, B) DO UPDATE")
841        );
842        assert!(!cleaned.to_uppercase().contains("WHERE A IS NOT NULL"));
843    }
844
845    #[test]
846    fn test_preprocess_postgres_preserves_unrelated_where() {
847        // The DELETE's WHERE is its own clause, not an ON-CONFLICT predicate;
848        // it must survive untouched.
849        let sql = "DELETE FROM t WHERE id = $1";
850        assert_eq!(preprocess_postgres_sql(sql), sql);
851    }
852
853    #[test]
854    fn test_preprocess_oracle_colon_placeholders() {
855        assert_eq!(
856            preprocess_oracle_sql("SELECT * FROM users WHERE id = :1"),
857            "SELECT * FROM users WHERE id = ?"
858        );
859        assert_eq!(
860            preprocess_oracle_sql("INSERT INTO users (name, email) VALUES (:1, :2)"),
861            "INSERT INTO users (name, email) VALUES (?, ?)"
862        );
863    }
864
865    #[test]
866    fn test_preprocess_oracle_preserves_string_literals() {
867        assert_eq!(
868            preprocess_oracle_sql("SELECT * FROM users WHERE name = ':1' AND id = :1"),
869            "SELECT * FROM users WHERE name = ':1' AND id = ?"
870        );
871    }
872
873    #[test]
874    fn test_preprocess_oracle_strips_returning_into() {
875        assert_eq!(
876            preprocess_oracle_sql(
877                "INSERT INTO users (name) VALUES (:1) RETURNING id, name INTO :2, :3"
878            ),
879            "INSERT INTO users (name) VALUES (?) RETURNING id, name"
880        );
881    }
882
883    #[test]
884    fn test_preprocess_oracle_full_insert_returning_into() {
885        let sql = "INSERT INTO users (name, email, active) VALUES (:1, :2, :3) RETURNING id, name, email, active, created_at INTO :4, :5, :6, :7, :8";
886        let result = preprocess_oracle_sql(sql);
887        assert_eq!(
888            result,
889            "INSERT INTO users (name, email, active) VALUES (?, ?, ?) RETURNING id, name, email, active, created_at"
890        );
891    }
892
893    #[test]
894    fn test_preprocess_oracle_no_returning_into_unchanged() {
895        assert_eq!(
896            preprocess_oracle_sql("DELETE FROM users WHERE id = :1"),
897            "DELETE FROM users WHERE id = ?"
898        );
899    }
900
901    #[test]
902    fn test_preprocess_mssql_single_placeholder() {
903        assert_eq!(
904            preprocess_mssql_sql("SELECT * FROM users WHERE id = @p1"),
905            "SELECT * FROM users WHERE id = ?"
906        );
907    }
908
909    #[test]
910    fn test_preprocess_mssql_multiple_placeholders() {
911        assert_eq!(
912            preprocess_mssql_sql("INSERT INTO users (name, email) VALUES (@p1, @p2)"),
913            "INSERT INTO users (name, email) VALUES (?, ?)"
914        );
915    }
916
917    #[test]
918    fn test_preprocess_mssql_preserves_string_literals() {
919        assert_eq!(
920            preprocess_mssql_sql("SELECT * FROM users WHERE name = '@p1' AND id = @p1"),
921            "SELECT * FROM users WHERE name = '@p1' AND id = ?"
922        );
923    }
924
925    #[test]
926    fn test_preprocess_mssql_case_insensitive_p() {
927        assert_eq!(
928            preprocess_mssql_sql("SELECT * FROM users WHERE id = @P1"),
929            "SELECT * FROM users WHERE id = ?"
930        );
931    }
932
933    #[test]
934    fn test_preprocess_mssql_non_placeholder_at_variable_unchanged() {
935        // @variable (not @pN pattern) must not be touched
936        assert_eq!(preprocess_mssql_sql("SELECT @myvar"), "SELECT @myvar");
937    }
938
939    #[test]
940    fn test_preprocess_mssql_multi_digit_placeholder() {
941        assert_eq!(preprocess_mssql_sql("SELECT @p10, @p2"), "SELECT ?, ?");
942    }
943
944    #[test]
945    fn test_preprocess_mssql_output_inserted_simple() {
946        let sql =
947            "INSERT INTO users (id, name) OUTPUT INSERTED.id, INSERTED.name VALUES (@p1, @p2)";
948        let result = preprocess_mssql_sql(sql);
949        // Should convert OUTPUT INSERTED.col to RETURNING col and @pN to ?
950        assert!(result.contains("RETURNING id, name"), "got: {}", result);
951        assert!(result.contains("VALUES (?, ?)"), "got: {}", result);
952        assert!(!result.contains("OUTPUT"), "got: {}", result);
953    }
954
955    #[test]
956    fn test_preprocess_mssql_output_inserted_full_example() {
957        let sql = "INSERT INTO users (id, name, email, active) OUTPUT INSERTED.id, INSERTED.name, INSERTED.email, INSERTED.active, INSERTED.created_at VALUES (@p1, @p2, @p3, @p4)";
958        let result = preprocess_mssql_sql(sql);
959        assert!(
960            result.contains("RETURNING id, name, email, active, created_at"),
961            "got: {}",
962            result
963        );
964        assert!(result.contains("VALUES (?, ?, ?, ?)"), "got: {}", result);
965    }
966
967    #[test]
968    fn test_preprocess_mssql_output_case_insensitive() {
969        let sql = "INSERT INTO users (id) output inserted.id values (@p1)";
970        let result = preprocess_mssql_sql(sql);
971        assert!(result.contains("RETURNING id"), "got: {}", result);
972        // The original lowercase "values" is preserved, then @p1 becomes ?
973        assert!(
974            result.contains("values (?)") || result.contains("VALUES (?)"),
975            "got: {}",
976            result
977        );
978    }
979
980    #[test]
981    fn test_preprocess_mssql_no_output_unchanged() {
982        let sql = "INSERT INTO users (id, name) VALUES (@p1, @p2)";
983        let result = preprocess_mssql_sql(sql);
984        assert_eq!(result, "INSERT INTO users (id, name) VALUES (?, ?)");
985    }
986
987    #[test]
988    fn test_preprocess_mssql_output_with_string_literal() {
989        // @p1 inside a string should be preserved by placeholder conversion
990        let sql =
991            "INSERT INTO users (id, name) OUTPUT INSERTED.id, INSERTED.name VALUES (@p1, '@p2')";
992        let result = preprocess_mssql_sql(sql);
993        assert!(result.contains("RETURNING id, name"), "got: {}", result);
994        assert!(result.contains("(?, '@p2')"), "got: {}", result);
995    }
996
997    #[test]
998    fn test_preprocess_mssql_output_with_whitespace() {
999        let sql =
1000            "INSERT INTO users (id, name)\nOUTPUT INSERTED.id,\n  INSERTED.name\nVALUES (@p1, @p2)";
1001        let result = preprocess_mssql_sql(sql);
1002        assert!(result.contains("RETURNING id, name"), "got: {}", result);
1003        assert!(result.contains("VALUES (?, ?)"), "got: {}", result);
1004    }
1005}