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    // Strip line comments + string literals first so we only scan structural SQL.
307    // (We still emit the original `sql` slice byte-for-byte; the upper-mask is
308    //  only used to decide *where* to cut.)
309    let mask = mask_postgres_for_scan(sql);
310    let mask_bytes = mask.as_bytes();
311    let bytes = sql.as_bytes();
312    let mut search_from = 0;
313    let mut result = String::with_capacity(sql.len());
314    let mut last = 0;
315    while let Some(rel) = find_keyword(&mask[search_from..], "ON CONFLICT") {
316        let on_conflict_pos = search_from + rel;
317        let after_on_conflict = on_conflict_pos + "ON CONFLICT".len();
318        let mut idx = after_on_conflict;
319        while idx < mask_bytes.len() && mask_bytes[idx].is_ascii_whitespace() {
320            idx += 1;
321        }
322        if idx >= mask_bytes.len() || mask_bytes[idx] != b'(' {
323            search_from = after_on_conflict;
324            continue;
325        }
326        let mut depth = 0i32;
327        let mut close = idx;
328        while close < mask_bytes.len() {
329            match mask_bytes[close] {
330                b'(' => depth += 1,
331                b')' => {
332                    depth -= 1;
333                    if depth == 0 {
334                        break;
335                    }
336                }
337                _ => {}
338            }
339            close += 1;
340        }
341        if depth != 0 {
342            return sql.to_string();
343        }
344        let mut after_cols = close + 1;
345        while after_cols < mask_bytes.len() && mask_bytes[after_cols].is_ascii_whitespace() {
346            after_cols += 1;
347        }
348        if mask[after_cols..].starts_with("WHERE")
349            && let Some(do_rel) = find_keyword(&mask[after_cols + "WHERE".len()..], "DO")
350        {
351            let do_abs = after_cols + "WHERE".len() + do_rel;
352            // Slice from the original SQL (preserves casing + UTF-8) up to
353            // the byte before WHERE; skip ahead to DO.
354            result.push_str(std::str::from_utf8(&bytes[last..after_cols]).unwrap_or(""));
355            last = do_abs;
356            search_from = do_abs;
357            continue;
358        }
359        search_from = close + 1;
360    }
361    result.push_str(std::str::from_utf8(&bytes[last..]).unwrap_or(""));
362    result
363}
364
365/// Build an ASCII-uppercase, fixed-byte-offset mask of `sql` where `--` line
366/// comments, `/* … */` block comments, and `'…'` / `$$…$$` string literals are
367/// replaced with spaces. Multi-byte UTF-8 is collapsed to ASCII spaces of the
368/// same byte length so positions in the mask line up with the original `sql`.
369fn mask_postgres_for_scan(sql: &str) -> String {
370    let bytes = sql.as_bytes();
371    let mut out = vec![b' '; bytes.len()];
372    let mut i = 0;
373    while i < bytes.len() {
374        let b = bytes[i];
375        if b == b'-' && i + 1 < bytes.len() && bytes[i + 1] == b'-' {
376            // Line comment — replace through end-of-line with spaces.
377            while i < bytes.len() && bytes[i] != b'\n' {
378                out[i] = b' ';
379                i += 1;
380            }
381            continue;
382        }
383        if b == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
384            // Block comment — replace through `*/`.
385            out[i] = b' ';
386            out[i + 1] = b' ';
387            i += 2;
388            while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
389                out[i] = b' ';
390                i += 1;
391            }
392            if i + 1 < bytes.len() {
393                out[i] = b' ';
394                out[i + 1] = b' ';
395                i += 2;
396            }
397            continue;
398        }
399        if b == b'\'' {
400            out[i] = b' ';
401            i += 1;
402            while i < bytes.len() {
403                if bytes[i] == b'\'' {
404                    if i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
405                        out[i] = b' ';
406                        out[i + 1] = b' ';
407                        i += 2;
408                        continue;
409                    }
410                    out[i] = b' ';
411                    i += 1;
412                    break;
413                }
414                out[i] = b' ';
415                i += 1;
416            }
417            continue;
418        }
419        // ASCII goes through as-uppercase; non-ASCII bytes become spaces so the
420        // mask stays single-byte-per-position and positions line up.
421        if b.is_ascii() {
422            out[i] = b.to_ascii_uppercase();
423        } else {
424            out[i] = b' ';
425        }
426        i += 1;
427    }
428    String::from_utf8(out).expect("mask is ASCII by construction")
429}
430
431/// Locate a whitespace-separated keyword in an uppercase haystack. Returns the
432/// byte offset of the keyword's start, or None if not found.
433fn find_keyword(haystack: &str, keyword: &str) -> Option<usize> {
434    let bytes = haystack.as_bytes();
435    let key = keyword.as_bytes();
436    let mut i = 0;
437    while i + key.len() <= bytes.len() {
438        if &bytes[i..i + key.len()] == key {
439            let prev_ok = i == 0 || !bytes[i - 1].is_ascii_alphanumeric();
440            let next = i + key.len();
441            let next_ok = next >= bytes.len() || !bytes[next].is_ascii_alphanumeric();
442            if prev_ok && next_ok {
443                return Some(i);
444            }
445        }
446        i += 1;
447    }
448    None
449}
450
451/// Preprocess Oracle SQL before parsing:
452/// 1. Strip `INTO :N, :N, ...` suffix from `RETURNING ... INTO` clauses.
453/// 2. Convert `:N` positional placeholders to `?` (universally supported).
454fn preprocess_oracle_sql(sql: &str) -> String {
455    // Strip Oracle RETURNING ... INTO clause (output bind variables)
456    // e.g. "INSERT ... RETURNING id, name INTO :4, :5" → "INSERT ... RETURNING id, name"
457    let sql = strip_returning_into(sql);
458
459    // Convert :N → ? (outside string literals)
460    let mut result = String::with_capacity(sql.len());
461    let mut chars = sql.chars().peekable();
462    while let Some(ch) = chars.next() {
463        if ch == '\'' {
464            // Skip string literals
465            result.push(ch);
466            while let Some(inner) = chars.next() {
467                result.push(inner);
468                if inner == '\'' {
469                    if chars.peek() == Some(&'\'') {
470                        result.push(chars.next().unwrap());
471                    } else {
472                        break;
473                    }
474                }
475            }
476        } else if ch == ':' && chars.peek().is_some_and(|c| c.is_ascii_digit()) {
477            // Convert :N → ?
478            result.push('?');
479            while chars.peek().is_some_and(|c| c.is_ascii_digit()) {
480                chars.next();
481            }
482        } else {
483            result.push(ch);
484        }
485    }
486    result
487}
488
489/// Convert MSSQL `@pN` positional placeholders to `?` (outside string literals).
490/// MsSqlDialect treats `@` as an identifier start, so `@p1` becomes an identifier
491/// rather than a `Placeholder` token — preprocessing normalises it to `?`.
492fn convert_mssql_placeholders(sql: &str) -> String {
493    let mut result = String::with_capacity(sql.len());
494    let mut chars = sql.chars().peekable();
495    while let Some(ch) = chars.next() {
496        if ch == '\'' {
497            // Skip string literals verbatim
498            result.push(ch);
499            while let Some(inner) = chars.next() {
500                result.push(inner);
501                if inner == '\'' {
502                    if chars.peek() == Some(&'\'') {
503                        // Escaped quote inside string literal
504                        result.push(chars.next().unwrap());
505                    } else {
506                        break;
507                    }
508                }
509            }
510        } else if ch == '@' && chars.peek().is_some_and(|c| *c == 'p' || *c == 'P') {
511            // Peek ahead: must be `@p` followed by at least one digit
512            let mut lookahead = chars.clone();
513            lookahead.next(); // consume the 'p'/'P'
514            if lookahead.peek().is_some_and(|c| c.is_ascii_digit()) {
515                // It is an `@pN` placeholder — consume `p` and all digits
516                chars.next(); // consume 'p'/'P'
517                while chars.peek().is_some_and(|c| c.is_ascii_digit()) {
518                    chars.next();
519                }
520                result.push('?');
521            } else {
522                result.push(ch);
523            }
524        } else {
525            result.push(ch);
526        }
527    }
528    result
529}
530
531/// Preprocess MSSQL SQL before parsing:
532/// 1. Strip `OUTPUT INSERTED.col, ...` clauses and convert to RETURNING
533/// 2. Convert `@pN` positional placeholders to `?`
534fn preprocess_mssql_sql(sql: &str) -> String {
535    // First pass: convert OUTPUT INSERTED.col to RETURNING col
536    let sql = strip_and_convert_mssql_output(sql);
537    // Second pass: convert @pN to ?
538    convert_mssql_placeholders(&sql)
539}
540
541/// Strip MSSQL `OUTPUT INSERTED.col1, INSERTED.col2, ...` from INSERT statements
542/// and convert it to a `RETURNING col1, col2, ...` clause.
543/// The OUTPUT clause appears between the column list and VALUES clause:
544///   INSERT INTO table (cols) OUTPUT INSERTED.col1, INSERTED.col2, ... VALUES (...)
545/// becomes:
546///   INSERT INTO table (cols) VALUES (...) RETURNING col1, col2, ...
547fn strip_and_convert_mssql_output(sql: &str) -> String {
548    // Case-insensitive search for OUTPUT keyword in INSERT statements
549    let upper = sql.to_uppercase();
550
551    // Only process INSERT statements with OUTPUT
552    if !upper.contains("INSERT") || !upper.contains("OUTPUT") {
553        return sql.to_string();
554    }
555
556    // Find the OUTPUT keyword
557    if let Some(output_pos) = find_word_position(&upper, "OUTPUT") {
558        // Check if this is actually part of an INSERT statement by finding INSERT before it
559        let before_output = &upper[..output_pos];
560        if !before_output.contains("INSERT") {
561            return sql.to_string();
562        }
563
564        // Look for the VALUES keyword after OUTPUT
565        let after_output = &upper[output_pos + "OUTPUT".len()..];
566        if let Some(values_offset) = find_word_position(after_output, "VALUES") {
567            let values_pos = output_pos + "OUTPUT".len() + values_offset;
568
569            // Extract the OUTPUT column list (between OUTPUT and VALUES)
570            let output_cols_str = &sql[output_pos + "OUTPUT".len()..values_pos];
571
572            // Parse column names: strip "INSERTED." prefix from each column name
573            let cols = parse_inserted_columns(output_cols_str);
574
575            if !cols.is_empty() {
576                // Build result: keep everything before OUTPUT, then VALUES clause,
577                // then RETURNING clause (before any trailing semicolon)
578                let before_output_sql = sql[..output_pos].trim_end();
579                let after_values = sql[values_pos..].trim_end();
580                let (values_body, trailing) = if let Some(stripped) = after_values.strip_suffix(';')
581                {
582                    (stripped, ";")
583                } else {
584                    (after_values, "")
585                };
586
587                return format!(
588                    "{}\n{} RETURNING {}{}",
589                    before_output_sql, values_body, cols, trailing
590                );
591            }
592        }
593    }
594
595    sql.to_string()
596}
597
598/// Find the position of a word (case-insensitive) in the text.
599/// The word must be a separate word, not part of another identifier.
600fn find_word_position(text: &str, word: &str) -> Option<usize> {
601    let mut pos = 0;
602    let word_len = word.len();
603    while let Some(idx) = text[pos..].find(word) {
604        let abs_idx = pos + idx;
605
606        // Check character before
607        let before_ok = abs_idx == 0
608            || !text
609                .as_bytes()
610                .get(abs_idx - 1)
611                .is_some_and(|&b| b.is_ascii_alphanumeric() || b == b'_');
612
613        // Check character after
614        let after_idx = abs_idx + word_len;
615        let after_ok = after_idx >= text.len()
616            || !text
617                .as_bytes()
618                .get(after_idx)
619                .is_some_and(|&b| b.is_ascii_alphanumeric() || b == b'_');
620
621        if before_ok && after_ok {
622            return Some(abs_idx);
623        }
624        pos = abs_idx + 1;
625    }
626    None
627}
628
629/// Parse INSERTED.col1, INSERTED.col2, ... and extract column names as "col1, col2, ..."
630fn parse_inserted_columns(output_str: &str) -> String {
631    let mut cols = Vec::new();
632
633    for part in output_str.split(',') {
634        let trimmed = part.trim();
635
636        // Try to extract column name after INSERTED.
637        if let Some(after_inserted) = trimmed
638            .strip_prefix("INSERTED.")
639            .or_else(|| trimmed.strip_prefix("inserted."))
640            .or_else(|| trimmed.strip_prefix("INSERTED"))
641            .or_else(|| trimmed.strip_prefix("inserted"))
642        {
643            let col_name = after_inserted.trim().to_string();
644            if !col_name.is_empty() {
645                cols.push(col_name);
646            }
647        }
648    }
649
650    cols.join(", ")
651}
652
653/// Strip the `INTO :N, :N, ...` suffix from an Oracle `RETURNING ... INTO` clause.
654fn strip_returning_into(sql: &str) -> String {
655    // Case-insensitive search for "INTO" after "RETURNING" at the end of the statement
656    let upper = sql.to_uppercase();
657    if let Some(ret_pos) = upper.rfind("RETURNING") {
658        let after_returning = &upper[ret_pos + "RETURNING".len()..];
659        if let Some(into_offset) = after_returning.find("INTO") {
660            let into_pos = ret_pos + "RETURNING".len() + into_offset;
661            // Keep everything before INTO, trim trailing whitespace/semicolons
662            let trimmed = sql[..into_pos].trim_end();
663            return trimmed.to_string();
664        }
665    }
666    sql.to_string()
667}
668
669#[cfg(test)]
670mod tests {
671    use super::*;
672    use crate::errors::ErrorCode;
673
674    fn parse(sql: &str) -> Result<Query, ScytheError> {
675        parse_query(sql)
676    }
677
678    #[test]
679    fn test_basic_parse() {
680        let input = "-- @name GetUsers\n-- @returns :many\nSELECT * FROM users;";
681        let q = parse(input).unwrap();
682        assert_eq!(q.name, "GetUsers");
683        assert_eq!(q.command, QueryCommand::Many);
684        assert!(q.sql.contains("SELECT"));
685    }
686
687    #[test]
688    fn test_all_command_types() {
689        let cases = vec![
690            (":one", QueryCommand::One),
691            (":many", QueryCommand::Many),
692            (":exec", QueryCommand::Exec),
693            (":exec_result", QueryCommand::ExecResult),
694            (":exec_rows", QueryCommand::ExecRows),
695        ];
696        for (tag, expected) in cases {
697            let input = format!("-- @name Q\n-- @returns {}\nSELECT 1", tag);
698            let q = parse(&input).unwrap();
699            assert_eq!(q.command, expected, "failed for {}", tag);
700        }
701    }
702
703    #[test]
704    fn test_case_insensitive_keywords() {
705        let input = "-- @Name GetUsers\n-- @RETURNS :many\nSELECT 1";
706        let q = parse(input).unwrap();
707        assert_eq!(q.name, "GetUsers");
708        assert_eq!(q.command, QueryCommand::Many);
709    }
710
711    #[test]
712    fn test_missing_name_errors() {
713        let input = "-- @returns :many\nSELECT 1";
714        let err = parse(input).unwrap_err();
715        assert_eq!(err.code, ErrorCode::MissingAnnotation);
716        assert!(err.message.contains("name"));
717    }
718
719    #[test]
720    fn test_missing_returns_errors() {
721        let input = "-- @name Foo\nSELECT 1";
722        let err = parse(input).unwrap_err();
723        assert_eq!(err.code, ErrorCode::MissingAnnotation);
724        assert!(err.message.contains("returns"));
725    }
726
727    #[test]
728    fn test_invalid_returns_value() {
729        let input = "-- @name Foo\n-- @returns :invalid\nSELECT 1";
730        let err = parse(input).unwrap_err();
731        assert_eq!(err.code, ErrorCode::InvalidAnnotation);
732    }
733
734    #[test]
735    fn test_empty_name_value() {
736        // An empty name is accepted by the parser (it stores "")
737        let input = "-- @name\n-- @returns :one\nSELECT 1";
738        let q = parse(input).unwrap();
739        assert_eq!(q.name, "");
740    }
741
742    #[test]
743    fn test_param_annotation() {
744        let input = "-- @name Foo\n-- @returns :one\n-- @param id: the user ID\nSELECT 1";
745        let q = parse(input).unwrap();
746        assert_eq!(q.annotations.param_docs.len(), 1);
747        assert_eq!(q.annotations.param_docs[0].name, "id");
748        assert_eq!(q.annotations.param_docs[0].description, "the user ID");
749    }
750
751    #[test]
752    fn test_param_no_description() {
753        let input = "-- @name Foo\n-- @returns :one\n-- @param id\nSELECT 1";
754        let q = parse(input).unwrap();
755        assert_eq!(q.annotations.param_docs.len(), 1);
756        assert_eq!(q.annotations.param_docs[0].name, "id");
757        assert_eq!(q.annotations.param_docs[0].description, "");
758    }
759
760    #[test]
761    fn test_nullable_annotation() {
762        let input = "-- @name Foo\n-- @returns :one\n-- @nullable col1, col2\nSELECT 1";
763        let q = parse(input).unwrap();
764        assert_eq!(q.annotations.nullable_overrides, vec!["col1", "col2"]);
765    }
766
767    #[test]
768    fn test_nonnull_annotation() {
769        let input = "-- @name Foo\n-- @returns :one\n-- @nonnull col1\nSELECT 1";
770        let q = parse(input).unwrap();
771        assert_eq!(q.annotations.nonnull_overrides, vec!["col1"]);
772    }
773
774    #[test]
775    fn test_json_annotation() {
776        let input = "-- @name Foo\n-- @returns :one\n-- @json data = EventData\nSELECT 1";
777        let q = parse(input).unwrap();
778        assert_eq!(q.annotations.json_mappings.len(), 1);
779        assert_eq!(q.annotations.json_mappings[0].column, "data");
780        assert_eq!(q.annotations.json_mappings[0].rust_type, "EventData");
781    }
782
783    #[test]
784    fn test_deprecated_annotation() {
785        let input = "-- @name Foo\n-- @returns :one\n-- @deprecated Use V2\nSELECT 1";
786        let q = parse(input).unwrap();
787        assert_eq!(q.annotations.deprecated, Some("Use V2".to_string()));
788    }
789
790    #[test]
791    fn test_sql_syntax_error() {
792        let input = "-- @name Foo\n-- @returns :one\nSELCT * FROM users";
793        let err = parse(input).unwrap_err();
794        assert_eq!(err.code, ErrorCode::SyntaxError);
795    }
796
797    #[test]
798    fn test_trailing_semicolon() {
799        let input = "-- @name Foo\n-- @returns :one\nSELECT 1;";
800        let q = parse(input).unwrap();
801        assert_eq!(q.name, "Foo");
802    }
803
804    #[test]
805    fn test_multiple_statements_error() {
806        let input = "-- @name Foo\n-- @returns :one\nSELECT 1; SELECT 2;";
807        let err = parse(input).unwrap_err();
808        assert_eq!(err.code, ErrorCode::SyntaxError);
809    }
810
811    #[test]
812    fn test_sql_preserved_without_annotations() {
813        let input = "-- @name Foo\n-- @returns :one\nSELECT id, name FROM users WHERE id = $1";
814        let q = parse(input).unwrap();
815        assert_eq!(q.sql, "SELECT id, name FROM users WHERE id = $1");
816    }
817
818    #[test]
819    fn test_returns_without_colon_prefix() {
820        let input = "-- @name Foo\n-- @returns many\nSELECT 1";
821        let q = parse(input).unwrap();
822        assert_eq!(q.command, QueryCommand::Many);
823    }
824
825    #[test]
826    fn test_batch_command() {
827        let input = "-- @name Foo\n-- @returns :batch\nSELECT 1";
828        let q = parse(input).unwrap();
829        assert_eq!(q.command, QueryCommand::Batch);
830    }
831
832    #[test]
833    fn test_grouped_command_with_group_by() {
834        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";
835        let q = parse(input).unwrap();
836        assert_eq!(q.command, QueryCommand::Grouped);
837        assert_eq!(q.annotations.group_by, Some("users.id".to_string()));
838    }
839
840    #[test]
841    fn test_grouped_command_without_group_by_errors() {
842        let input = "-- @name Foo\n-- @returns :grouped\nSELECT 1";
843        let err = parse(input).unwrap_err();
844        assert_eq!(err.code, ErrorCode::InvalidAnnotation);
845        assert!(err.message.contains("@group_by"));
846    }
847
848    #[test]
849    fn test_group_by_without_grouped_is_ignored() {
850        let input = "-- @name Foo\n-- @returns :many\n-- @group_by users.id\nSELECT 1";
851        let q = parse(input).unwrap();
852        assert_eq!(q.command, QueryCommand::Many);
853        assert_eq!(q.annotations.group_by, Some("users.id".to_string()));
854    }
855
856    #[test]
857    fn test_preprocess_postgres_strips_partial_index_where() {
858        let sql = "INSERT INTO billing_events (project_id, stripe_event_id) \
859                   VALUES ($1, $2) \
860                   ON CONFLICT (stripe_event_id) WHERE stripe_event_id IS NOT NULL DO NOTHING";
861        let cleaned = preprocess_postgres_sql(sql);
862        assert!(
863            !cleaned
864                .to_uppercase()
865                .contains("WHERE STRIPE_EVENT_ID IS NOT NULL"),
866            "WHERE clause must be stripped between ON CONFLICT cols and DO; got: {cleaned}"
867        );
868        assert!(
869            cleaned
870                .to_uppercase()
871                .contains("ON CONFLICT (STRIPE_EVENT_ID) DO NOTHING")
872        );
873        // sqlparser must accept the cleaned form.
874        sqlparser::parser::Parser::parse_sql(&sqlparser::dialect::PostgreSqlDialect {}, &cleaned)
875            .expect("cleaned SQL should parse");
876    }
877
878    #[test]
879    fn test_preprocess_postgres_no_op_when_no_partial_clause() {
880        let sql = "INSERT INTO t (a) VALUES ($1) ON CONFLICT (a) DO UPDATE SET a = EXCLUDED.a";
881        assert_eq!(preprocess_postgres_sql(sql), sql);
882    }
883
884    #[test]
885    fn test_preprocess_postgres_leaves_on_conflict_on_constraint_alone() {
886        let sql = "INSERT INTO t (a) VALUES ($1) ON CONFLICT ON CONSTRAINT t_a_uidx DO NOTHING";
887        assert_eq!(preprocess_postgres_sql(sql), sql);
888    }
889
890    #[test]
891    fn test_preprocess_postgres_handles_compound_index_cols() {
892        let sql = "INSERT INTO t (a, b) VALUES ($1, $2) \
893                   ON CONFLICT (a, b) WHERE a IS NOT NULL AND b > 0 DO UPDATE SET b = EXCLUDED.b";
894        let cleaned = preprocess_postgres_sql(sql);
895        assert!(
896            cleaned
897                .to_uppercase()
898                .contains("ON CONFLICT (A, B) DO UPDATE")
899        );
900        assert!(!cleaned.to_uppercase().contains("WHERE A IS NOT NULL"));
901    }
902
903    #[test]
904    fn test_preprocess_postgres_preserves_unrelated_where() {
905        // The DELETE's WHERE is its own clause, not an ON-CONFLICT predicate;
906        // it must survive untouched.
907        let sql = "DELETE FROM t WHERE id = $1";
908        assert_eq!(preprocess_postgres_sql(sql), sql);
909    }
910
911    #[test]
912    fn test_preprocess_postgres_ignores_text_inside_line_comments() {
913        // Earlier scans treated this as a real `ON CONFLICT (col) WHERE … DO`
914        // and excised the entire comment + INSERT body up to the next `DO`.
915        // Comments must be opaque to the predicate-stripping pass.
916        let sql = "-- inline doc: `ON CONFLICT (col) WHERE …` is the partial form\n\
917                   INSERT INTO t (a) VALUES ($1) \
918                   ON CONFLICT (a) WHERE a IS NOT NULL DO NOTHING";
919        let cleaned = preprocess_postgres_sql(sql);
920        assert!(
921            cleaned.contains("-- inline doc"),
922            "comment must survive the pass; got: {cleaned}"
923        );
924        assert!(cleaned.contains("ON CONFLICT (a) DO NOTHING"));
925    }
926
927    #[test]
928    fn test_preprocess_postgres_ignores_text_inside_string_literals() {
929        let sql = "SELECT 'ON CONFLICT (a) WHERE a IS NOT NULL DO NOTHING' AS s";
930        assert_eq!(preprocess_postgres_sql(sql), sql);
931    }
932
933    #[test]
934    fn test_preprocess_oracle_colon_placeholders() {
935        assert_eq!(
936            preprocess_oracle_sql("SELECT * FROM users WHERE id = :1"),
937            "SELECT * FROM users WHERE id = ?"
938        );
939        assert_eq!(
940            preprocess_oracle_sql("INSERT INTO users (name, email) VALUES (:1, :2)"),
941            "INSERT INTO users (name, email) VALUES (?, ?)"
942        );
943    }
944
945    #[test]
946    fn test_preprocess_oracle_preserves_string_literals() {
947        assert_eq!(
948            preprocess_oracle_sql("SELECT * FROM users WHERE name = ':1' AND id = :1"),
949            "SELECT * FROM users WHERE name = ':1' AND id = ?"
950        );
951    }
952
953    #[test]
954    fn test_preprocess_oracle_strips_returning_into() {
955        assert_eq!(
956            preprocess_oracle_sql(
957                "INSERT INTO users (name) VALUES (:1) RETURNING id, name INTO :2, :3"
958            ),
959            "INSERT INTO users (name) VALUES (?) RETURNING id, name"
960        );
961    }
962
963    #[test]
964    fn test_preprocess_oracle_full_insert_returning_into() {
965        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";
966        let result = preprocess_oracle_sql(sql);
967        assert_eq!(
968            result,
969            "INSERT INTO users (name, email, active) VALUES (?, ?, ?) RETURNING id, name, email, active, created_at"
970        );
971    }
972
973    #[test]
974    fn test_preprocess_oracle_no_returning_into_unchanged() {
975        assert_eq!(
976            preprocess_oracle_sql("DELETE FROM users WHERE id = :1"),
977            "DELETE FROM users WHERE id = ?"
978        );
979    }
980
981    #[test]
982    fn test_preprocess_mssql_single_placeholder() {
983        assert_eq!(
984            preprocess_mssql_sql("SELECT * FROM users WHERE id = @p1"),
985            "SELECT * FROM users WHERE id = ?"
986        );
987    }
988
989    #[test]
990    fn test_preprocess_mssql_multiple_placeholders() {
991        assert_eq!(
992            preprocess_mssql_sql("INSERT INTO users (name, email) VALUES (@p1, @p2)"),
993            "INSERT INTO users (name, email) VALUES (?, ?)"
994        );
995    }
996
997    #[test]
998    fn test_preprocess_mssql_preserves_string_literals() {
999        assert_eq!(
1000            preprocess_mssql_sql("SELECT * FROM users WHERE name = '@p1' AND id = @p1"),
1001            "SELECT * FROM users WHERE name = '@p1' AND id = ?"
1002        );
1003    }
1004
1005    #[test]
1006    fn test_preprocess_mssql_case_insensitive_p() {
1007        assert_eq!(
1008            preprocess_mssql_sql("SELECT * FROM users WHERE id = @P1"),
1009            "SELECT * FROM users WHERE id = ?"
1010        );
1011    }
1012
1013    #[test]
1014    fn test_preprocess_mssql_non_placeholder_at_variable_unchanged() {
1015        // @variable (not @pN pattern) must not be touched
1016        assert_eq!(preprocess_mssql_sql("SELECT @myvar"), "SELECT @myvar");
1017    }
1018
1019    #[test]
1020    fn test_preprocess_mssql_multi_digit_placeholder() {
1021        assert_eq!(preprocess_mssql_sql("SELECT @p10, @p2"), "SELECT ?, ?");
1022    }
1023
1024    #[test]
1025    fn test_preprocess_mssql_output_inserted_simple() {
1026        let sql =
1027            "INSERT INTO users (id, name) OUTPUT INSERTED.id, INSERTED.name VALUES (@p1, @p2)";
1028        let result = preprocess_mssql_sql(sql);
1029        // Should convert OUTPUT INSERTED.col to RETURNING col and @pN to ?
1030        assert!(result.contains("RETURNING id, name"), "got: {}", result);
1031        assert!(result.contains("VALUES (?, ?)"), "got: {}", result);
1032        assert!(!result.contains("OUTPUT"), "got: {}", result);
1033    }
1034
1035    #[test]
1036    fn test_preprocess_mssql_output_inserted_full_example() {
1037        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)";
1038        let result = preprocess_mssql_sql(sql);
1039        assert!(
1040            result.contains("RETURNING id, name, email, active, created_at"),
1041            "got: {}",
1042            result
1043        );
1044        assert!(result.contains("VALUES (?, ?, ?, ?)"), "got: {}", result);
1045    }
1046
1047    #[test]
1048    fn test_preprocess_mssql_output_case_insensitive() {
1049        let sql = "INSERT INTO users (id) output inserted.id values (@p1)";
1050        let result = preprocess_mssql_sql(sql);
1051        assert!(result.contains("RETURNING id"), "got: {}", result);
1052        // The original lowercase "values" is preserved, then @p1 becomes ?
1053        assert!(
1054            result.contains("values (?)") || result.contains("VALUES (?)"),
1055            "got: {}",
1056            result
1057        );
1058    }
1059
1060    #[test]
1061    fn test_preprocess_mssql_no_output_unchanged() {
1062        let sql = "INSERT INTO users (id, name) VALUES (@p1, @p2)";
1063        let result = preprocess_mssql_sql(sql);
1064        assert_eq!(result, "INSERT INTO users (id, name) VALUES (?, ?)");
1065    }
1066
1067    #[test]
1068    fn test_preprocess_mssql_output_with_string_literal() {
1069        // @p1 inside a string should be preserved by placeholder conversion
1070        let sql =
1071            "INSERT INTO users (id, name) OUTPUT INSERTED.id, INSERTED.name VALUES (@p1, '@p2')";
1072        let result = preprocess_mssql_sql(sql);
1073        assert!(result.contains("RETURNING id, name"), "got: {}", result);
1074        assert!(result.contains("(?, '@p2')"), "got: {}", result);
1075    }
1076
1077    #[test]
1078    fn test_preprocess_mssql_output_with_whitespace() {
1079        let sql =
1080            "INSERT INTO users (id, name)\nOUTPUT INSERTED.id,\n  INSERTED.name\nVALUES (@p1, @p2)";
1081        let result = preprocess_mssql_sql(sql);
1082        assert!(result.contains("RETURNING id, name"), "got: {}", result);
1083        assert!(result.contains("VALUES (?, ?)"), "got: {}", result);
1084    }
1085}
scythe_core/parser/mod.rs

scythe_core/parser/
mod.rs
