scythe_core/parser/
mod.rs

1use sqlparser::parser::Parser;
2
3use crate::dialect::SqlDialect;
4use crate::errors::ScytheError;
5
6#[derive(Debug, Clone, Default, PartialEq, Eq)]
7#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
8pub enum QueryCommand {
9    One,
10    Opt,
11    Many,
12    #[default]
13    Exec,
14    ExecResult,
15    ExecRows,
16    Batch,
17    Grouped,
18}
19
20impl std::fmt::Display for QueryCommand {
21    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
22        match self {
23            QueryCommand::One => write!(f, "one"),
24            QueryCommand::Opt => write!(f, "opt"),
25            QueryCommand::Many => write!(f, "many"),
26            QueryCommand::Exec => write!(f, "exec"),
27            QueryCommand::ExecResult => write!(f, "exec_result"),
28            QueryCommand::ExecRows => write!(f, "exec_rows"),
29            QueryCommand::Batch => write!(f, "batch"),
30            QueryCommand::Grouped => write!(f, "grouped"),
31        }
32    }
33}
34
35impl QueryCommand {
36    fn from_str(s: &str) -> Result<Self, ScytheError> {
37        match s {
38            "one" => Ok(QueryCommand::One),
39            "opt" => Ok(QueryCommand::Opt),
40            "many" => Ok(QueryCommand::Many),
41            "exec" => Ok(QueryCommand::Exec),
42            "exec_result" => Ok(QueryCommand::ExecResult),
43            "exec_rows" => Ok(QueryCommand::ExecRows),
44            "batch" => Ok(QueryCommand::Batch),
45            "grouped" => Ok(QueryCommand::Grouped),
46            other => Err(ScytheError::invalid_annotation(format!(
47                "invalid @returns value: {other}"
48            ))),
49        }
50    }
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
55pub struct ParamDoc {
56    pub name: String,
57    pub description: String,
58}
59
60/// An explicit `-- @param $N name[: description]` annotation that binds a
61/// position number to a human-chosen parameter name, overriding the `pN`
62/// fallback name that the analyzer would otherwise infer.
63///
64/// The positional form (`$N` as first token) is distinct from the existing
65/// docs-only `-- @param name: description` form, which is preserved unchanged
66/// for backward compatibility.
67#[derive(Debug, Clone, PartialEq, Eq)]
68#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
69pub struct PositionalParamDoc {
70    /// 1-based parameter position (`$N`).
71    pub position: i64,
72    /// Caller-chosen name that replaces the `pN` fallback.
73    pub name: String,
74    /// Optional description (everything after the `:` separator).
75    pub description: String,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq)]
79#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
80pub struct JsonMapping {
81    pub column: String,
82    pub rust_type: String,
83}
84
85/// A custom (non-native) annotation captured verbatim from the SQL source.
86///
87/// Scythe parses its known annotations (`@name`, `@returns`, `@param`, `@nullable`,
88/// `@nonnull`, `@json`, `@optional`, `@group_by`, `@deprecated`) into typed fields.
89/// Any other `-- @<name> <value>` line is captured here as an opaque triple and
90/// exposed to crate consumers, who can layer their own annotation vocabulary on
91/// top of scythe without coupling scythe to their domain.
92#[derive(Debug, Clone, PartialEq, Eq)]
93#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
94pub struct CustomAnnotation {
95    /// Annotation name, lowercased, without the leading `@` (e.g. `http`, `http_param`).
96    pub name: String,
97    /// Everything after the name on the line, trimmed. Empty if the annotation had no value.
98    pub value: String,
99    /// 1-based line number within the query SQL, for diagnostics.
100    pub line: usize,
101}
102
103#[derive(Debug, Clone, Default, PartialEq, Eq)]
104#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
105pub struct Annotations {
106    pub name: String,
107    pub command: QueryCommand,
108    pub param_docs: Vec<ParamDoc>,
109    pub nullable_overrides: Vec<String>,
110    pub nonnull_overrides: Vec<String>,
111    pub json_mappings: Vec<JsonMapping>,
112    pub deprecated: Option<String>,
113    pub optional_params: Vec<String>,
114    pub group_by: Option<String>,
115    /// Explicit `-- @param $N name[: description]` annotations that override
116    /// the inferred or fallback `pN` name for a specific positional parameter.
117    #[cfg_attr(feature = "serde", serde(default))]
118    pub positional_param_docs: Vec<PositionalParamDoc>,
119    /// Annotations scythe does not natively recognise, preserved in source order
120    /// for crate consumers to interpret.
121    pub custom: Vec<CustomAnnotation>,
122}
123
124#[derive(Debug)]
125pub struct Query {
126    pub name: String,
127    pub command: QueryCommand,
128    pub sql: String,
129    pub stmt: sqlparser::ast::Statement,
130    pub annotations: Annotations,
131}
132
133/// Parse a single annotated SQL query into a `Query` using the PostgreSQL dialect.
134pub fn parse_query(query_sql: &str) -> Result<Query, ScytheError> {
135    parse_query_with_dialect(query_sql, &SqlDialect::PostgreSQL)
136}
137
138/// Parse a single annotated SQL query into a `Query` using the specified dialect.
139pub fn parse_query_with_dialect(query_sql: &str, dialect: &SqlDialect) -> Result<Query, ScytheError> {
140    let mut name: Option<String> = None;
141    let mut command: Option<QueryCommand> = None;
142    let mut param_docs = Vec::new();
143    let mut positional_param_docs: Vec<PositionalParamDoc> = Vec::new();
144    let mut nullable_overrides = Vec::new();
145    let mut nonnull_overrides = Vec::new();
146    let mut json_mappings = Vec::new();
147    let mut deprecated: Option<String> = None;
148    let mut optional_params = Vec::new();
149    let mut group_by: Option<String> = None;
150    let mut custom: Vec<CustomAnnotation> = Vec::new();
151
152    let mut sql_lines = Vec::new();
153
154    for (line_idx, line) in query_sql.lines().enumerate() {
155        let line_no = line_idx + 1;
156        let trimmed = line.trim();
157
158        let annotation_body = if let Some(rest) = trimmed.strip_prefix("--") {
159            let rest = rest.trim_start();
160            rest.strip_prefix('@')
161        } else {
162            None
163        };
164
165        if let Some(body) = annotation_body {
166            let (keyword, value) = match body.find(|c: char| c.is_whitespace()) {
167                Some(pos) => (&body[..pos], body[pos..].trim()),
168                None => (body, ""),
169            };
170
171            match keyword.to_ascii_lowercase().as_str() {
172                "name" => {
173                    name = Some(value.to_string());
174                }
175                "returns" => {
176                    let cmd_str = value.strip_prefix(':').unwrap_or(value);
177                    command = Some(QueryCommand::from_str(cmd_str)?);
178                }
179                "param" => {
180                    //   Positional: -- @param $N name[: description]
181                    //   Docs-only:  -- @param name[: description]
182                    let first_end = value.find(|c: char| c.is_whitespace());
183                    let first_token = first_end.map(|p| &value[..p]).unwrap_or(value);
184
185                    if let Some(digits) = first_token.strip_prefix('$')
186                        && let Ok(pos) = digits.parse::<i64>()
187                        && pos > 0
188                    {
189                        let rest = first_end.map(|p| value[p..].trim()).unwrap_or("").trim();
190                        if !rest.is_empty() {
191                            let (param_name, description) = if let Some(colon_pos) = rest.find(':') {
192                                (
193                                    rest[..colon_pos].trim().to_string(),
194                                    rest[colon_pos + 1..].trim().to_string(),
195                                )
196                            } else {
197                                (rest.to_string(), String::new())
198                            };
199                            if !param_name.is_empty() {
200                                positional_param_docs.push(PositionalParamDoc {
201                                    position: pos,
202                                    name: param_name,
203                                    description,
204                                });
205                            }
206                        }
207                    } else {
208                        if let Some(colon_pos) = value.find(':') {
209                            let param_name = value[..colon_pos].trim().to_string();
210                            let description = value[colon_pos + 1..].trim().to_string();
211                            param_docs.push(ParamDoc {
212                                name: param_name,
213                                description,
214                            });
215                        } else {
216                            param_docs.push(ParamDoc {
217                                name: value.to_string(),
218                                description: String::new(),
219                            });
220                        }
221                    }
222                }
223                "nullable" => {
224                    for col in value.split(',') {
225                        let col = col.trim();
226                        if !col.is_empty() {
227                            nullable_overrides.push(col.to_string());
228                        }
229                    }
230                }
231                "nonnull" => {
232                    for col in value.split(',') {
233                        let col = col.trim();
234                        if !col.is_empty() {
235                            nonnull_overrides.push(col.to_string());
236                        }
237                    }
238                }
239                "json" => {
240                    if let Some(eq_pos) = value.find('=') {
241                        let column = value[..eq_pos].trim().to_string();
242                        let rust_type = value[eq_pos + 1..].trim().to_string();
243                        json_mappings.push(JsonMapping { column, rust_type });
244                    }
245                }
246                "deprecated" => {
247                    deprecated = Some(value.to_string());
248                }
249                "group_by" => {
250                    group_by = Some(value.to_string());
251                }
252                "optional" => {
253                    for param in value.split(',') {
254                        let param = param.trim();
255                        if !param.is_empty() {
256                            optional_params.push(param.to_string());
257                        }
258                    }
259                }
260                other => {
261                    custom.push(CustomAnnotation {
262                        name: other.to_string(),
263                        value: value.to_string(),
264                        line: line_no,
265                    });
266                }
267            }
268        } else {
269            sql_lines.push(line);
270        }
271    }
272
273    let name = name.ok_or_else(|| ScytheError::missing_annotation("name"))?;
274    let command = command.ok_or_else(|| ScytheError::missing_annotation("returns"))?;
275
276    if command == QueryCommand::Grouped && group_by.is_none() {
277        return Err(ScytheError::invalid_annotation(
278            "@returns :grouped requires a @group_by annotation (e.g. @group_by users.id)",
279        ));
280    }
281
282    let sql = sql_lines.join("\n").trim().to_string();
283
284    if sql.is_empty() {
285        return Err(ScytheError::syntax("empty SQL body"));
286    }
287
288    let (sql, parse_sql) = if *dialect == SqlDialect::Oracle {
289        let processed = preprocess_oracle_sql(&sql);
290        (processed.clone(), processed)
291    } else if *dialect == SqlDialect::MsSql {
292        let codegen_sql = convert_mssql_placeholders(&sql);
293        let parse_sql = preprocess_mssql_sql(&sql);
294        (codegen_sql, parse_sql)
295    } else if *dialect == SqlDialect::PostgreSQL {
296        let parse_sql = preprocess_postgres_sql(&sql);
297        (sql.clone(), parse_sql)
298    } else {
299        (sql.clone(), sql)
300    };
301
302    let parser_dialect = dialect.to_sqlparser_dialect();
303    let statements = Parser::parse_sql(parser_dialect.as_ref(), &parse_sql)
304        .map_err(|e| ScytheError::syntax(format!("syntax error: {}", e)))?;
305
306    if statements.len() != 1 {
307        let non_empty: Vec<_> = statements
308            .into_iter()
309            .filter(|s| !matches!(s, sqlparser::ast::Statement::Flush { .. }) && format!("{s}") != "")
310            .collect();
311        if non_empty.len() != 1 {
312            return Err(ScytheError::syntax("expected exactly one SQL statement"));
313        }
314        let stmt = non_empty.into_iter().next().expect("filtered to exactly one statement");
315        let annotations = Annotations {
316            name: name.clone(),
317            command: command.clone(),
318            param_docs,
319            positional_param_docs: positional_param_docs.clone(),
320            nullable_overrides,
321            nonnull_overrides,
322            json_mappings,
323            deprecated,
324            optional_params,
325            group_by: group_by.clone(),
326            custom,
327        };
328        return Ok(Query {
329            name,
330            command,
331            sql,
332            stmt,
333            annotations,
334        });
335    }
336
337    let stmt = statements
338        .into_iter()
339        .next()
340        .expect("filtered to exactly one statement");
341
342    let annotations = Annotations {
343        name: name.clone(),
344        command: command.clone(),
345        param_docs,
346        positional_param_docs,
347        nullable_overrides,
348        nonnull_overrides,
349        json_mappings,
350        deprecated,
351        optional_params,
352        group_by,
353        custom,
354    };
355
356    Ok(Query {
357        name,
358        command,
359        sql,
360        stmt,
361        annotations,
362    })
363}
364
365/// Strip the `WHERE …` predicate that PostgreSQL allows between
366/// `ON CONFLICT (cols)` and `DO …` (the index-inference form for partial
367/// unique indexes). sqlparser-rs through 0.61 does not parse this construct;
368/// we lift it out for the parser and let the caller keep the original SQL
369/// for codegen + runtime, where Postgres validates it.
370fn preprocess_postgres_sql(sql: &str) -> String {
371    let mask = mask_postgres_for_scan(sql);
372    let mask_bytes = mask.as_bytes();
373    let bytes = sql.as_bytes();
374    let mut search_from = 0;
375    let mut result = String::with_capacity(sql.len());
376    let mut last = 0;
377    while let Some(rel) = find_keyword(&mask[search_from..], "ON CONFLICT") {
378        let on_conflict_pos = search_from + rel;
379        let after_on_conflict = on_conflict_pos + "ON CONFLICT".len();
380        let mut idx = after_on_conflict;
381        while idx < mask_bytes.len() && mask_bytes[idx].is_ascii_whitespace() {
382            idx += 1;
383        }
384        if idx >= mask_bytes.len() || mask_bytes[idx] != b'(' {
385            search_from = after_on_conflict;
386            continue;
387        }
388        let mut depth = 0i32;
389        let mut close = idx;
390        while close < mask_bytes.len() {
391            match mask_bytes[close] {
392                b'(' => depth += 1,
393                b')' => {
394                    depth -= 1;
395                    if depth == 0 {
396                        break;
397                    }
398                }
399                _ => {}
400            }
401            close += 1;
402        }
403        if depth != 0 {
404            return sql.to_string();
405        }
406        let mut after_cols = close + 1;
407        while after_cols < mask_bytes.len() && mask_bytes[after_cols].is_ascii_whitespace() {
408            after_cols += 1;
409        }
410        if mask[after_cols..].starts_with("WHERE")
411            && let Some(do_rel) = find_keyword(&mask[after_cols + "WHERE".len()..], "DO")
412        {
413            let do_abs = after_cols + "WHERE".len() + do_rel;
414            result.push_str(std::str::from_utf8(&bytes[last..after_cols]).unwrap_or(""));
415            last = do_abs;
416            search_from = do_abs;
417            continue;
418        }
419        search_from = close + 1;
420    }
421    result.push_str(std::str::from_utf8(&bytes[last..]).unwrap_or(""));
422    result
423}
424
425/// Build an ASCII-uppercase, fixed-byte-offset mask of `sql` where `--` line
426/// comments, `/* … */` block comments, and `'…'` / `$$…$$` string literals are
427/// replaced with spaces. Multi-byte UTF-8 is collapsed to ASCII spaces of the
428/// same byte length so positions in the mask line up with the original `sql`.
429fn mask_postgres_for_scan(sql: &str) -> String {
430    let bytes = sql.as_bytes();
431    let mut out = vec![b' '; bytes.len()];
432    let mut i = 0;
433    while i < bytes.len() {
434        let b = bytes[i];
435        if b == b'-' && i + 1 < bytes.len() && bytes[i + 1] == b'-' {
436            while i < bytes.len() && bytes[i] != b'\n' {
437                out[i] = b' ';
438                i += 1;
439            }
440            continue;
441        }
442        if b == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
443            out[i] = b' ';
444            out[i + 1] = b' ';
445            i += 2;
446            while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
447                out[i] = b' ';
448                i += 1;
449            }
450            if i + 1 < bytes.len() {
451                out[i] = b' ';
452                out[i + 1] = b' ';
453                i += 2;
454            }
455            continue;
456        }
457        if b == b'\'' {
458            out[i] = b' ';
459            i += 1;
460            while i < bytes.len() {
461                if bytes[i] == b'\'' {
462                    if i + 1 < bytes.len() && bytes[i + 1] == b'\'' {
463                        out[i] = b' ';
464                        out[i + 1] = b' ';
465                        i += 2;
466                        continue;
467                    }
468                    out[i] = b' ';
469                    i += 1;
470                    break;
471                }
472                out[i] = b' ';
473                i += 1;
474            }
475            continue;
476        }
477        if b.is_ascii() {
478            out[i] = b.to_ascii_uppercase();
479        } else {
480            out[i] = b' ';
481        }
482        i += 1;
483    }
484    String::from_utf8(out).expect("mask is ASCII by construction")
485}
486
487/// Locate a whitespace-separated keyword in an uppercase haystack. Returns the
488/// byte offset of the keyword's start, or None if not found.
489fn find_keyword(haystack: &str, keyword: &str) -> Option<usize> {
490    let bytes = haystack.as_bytes();
491    let key = keyword.as_bytes();
492    let mut i = 0;
493    while i + key.len() <= bytes.len() {
494        if &bytes[i..i + key.len()] == key {
495            let prev_ok = i == 0 || !bytes[i - 1].is_ascii_alphanumeric();
496            let next = i + key.len();
497            let next_ok = next >= bytes.len() || !bytes[next].is_ascii_alphanumeric();
498            if prev_ok && next_ok {
499                return Some(i);
500            }
501        }
502        i += 1;
503    }
504    None
505}
506
507/// Preprocess Oracle SQL before parsing:
508/// 1. Strip `INTO :N, :N, ...` suffix from `RETURNING ... INTO` clauses.
509/// 2. Convert `:N` positional placeholders to `?` (universally supported).
510fn preprocess_oracle_sql(sql: &str) -> String {
511    let sql = strip_returning_into(sql);
512
513    let mut result = String::with_capacity(sql.len());
514    let mut chars = sql.chars().peekable();
515    while let Some(ch) = chars.next() {
516        if ch == '\'' {
517            result.push(ch);
518            while let Some(inner) = chars.next() {
519                result.push(inner);
520                if inner == '\'' {
521                    if chars.peek() == Some(&'\'') {
522                        result.push(chars.next().unwrap());
523                    } else {
524                        break;
525                    }
526                }
527            }
528        } else if ch == ':' && chars.peek().is_some_and(|c| c.is_ascii_digit()) {
529            result.push('?');
530            while chars.peek().is_some_and(|c| c.is_ascii_digit()) {
531                chars.next();
532            }
533        } else {
534            result.push(ch);
535        }
536    }
537    result
538}
539
540/// Convert MSSQL `@pN` positional placeholders to `?` (outside string literals).
541/// MsSqlDialect treats `@` as an identifier start, so `@p1` becomes an identifier
542/// rather than a `Placeholder` token — preprocessing normalises it to `?`.
543fn convert_mssql_placeholders(sql: &str) -> String {
544    let mut result = String::with_capacity(sql.len());
545    let mut chars = sql.chars().peekable();
546    while let Some(ch) = chars.next() {
547        if ch == '\'' {
548            result.push(ch);
549            while let Some(inner) = chars.next() {
550                result.push(inner);
551                if inner == '\'' {
552                    if chars.peek() == Some(&'\'') {
553                        result.push(chars.next().unwrap());
554                    } else {
555                        break;
556                    }
557                }
558            }
559        } else if ch == '@' && chars.peek().is_some_and(|c| *c == 'p' || *c == 'P') {
560            let mut lookahead = chars.clone();
561            lookahead.next();
562            if lookahead.peek().is_some_and(|c| c.is_ascii_digit()) {
563                chars.next();
564                while chars.peek().is_some_and(|c| c.is_ascii_digit()) {
565                    chars.next();
566                }
567                result.push('?');
568            } else {
569                result.push(ch);
570            }
571        } else {
572            result.push(ch);
573        }
574    }
575    result
576}
577
578/// Preprocess MSSQL SQL before parsing:
579/// 1. Strip `OUTPUT INSERTED.col, ...` clauses and convert to RETURNING
580/// 2. Convert `@pN` positional placeholders to `?`
581fn preprocess_mssql_sql(sql: &str) -> String {
582    let sql = strip_and_convert_mssql_output(sql);
583    convert_mssql_placeholders(&sql)
584}
585
586/// Strip MSSQL `OUTPUT INSERTED.col1, INSERTED.col2, ...` from INSERT statements
587/// and convert it to a `RETURNING col1, col2, ...` clause.
588/// The OUTPUT clause appears between the column list and VALUES clause:
589///   INSERT INTO table (cols) OUTPUT INSERTED.col1, INSERTED.col2, ... VALUES (...)
590/// becomes:
591///   INSERT INTO table (cols) VALUES (...) RETURNING col1, col2, ...
592fn strip_and_convert_mssql_output(sql: &str) -> String {
593    let upper = sql.to_uppercase();
594
595    if !upper.contains("INSERT") || !upper.contains("OUTPUT") {
596        return sql.to_string();
597    }
598
599    if let Some(output_pos) = find_word_position(&upper, "OUTPUT") {
600        let before_output = &upper[..output_pos];
601        if !before_output.contains("INSERT") {
602            return sql.to_string();
603        }
604
605        let after_output = &upper[output_pos + "OUTPUT".len()..];
606        if let Some(values_offset) = find_word_position(after_output, "VALUES") {
607            let values_pos = output_pos + "OUTPUT".len() + values_offset;
608
609            let output_cols_str = &sql[output_pos + "OUTPUT".len()..values_pos];
610
611            let cols = parse_inserted_columns(output_cols_str);
612
613            if !cols.is_empty() {
614                let before_output_sql = sql[..output_pos].trim_end();
615                let after_values = sql[values_pos..].trim_end();
616                let (values_body, trailing) = if let Some(stripped) = after_values.strip_suffix(';') {
617                    (stripped, ";")
618                } else {
619                    (after_values, "")
620                };
621
622                return format!("{}\n{} RETURNING {}{}", before_output_sql, values_body, cols, trailing);
623            }
624        }
625    }
626
627    sql.to_string()
628}
629
630/// Find the position of a word (case-insensitive) in the text.
631/// The word must be a separate word, not part of another identifier.
632fn find_word_position(text: &str, word: &str) -> Option<usize> {
633    let mut pos = 0;
634    let word_len = word.len();
635    while let Some(idx) = text[pos..].find(word) {
636        let abs_idx = pos + idx;
637
638        let before_ok = abs_idx == 0
639            || !text
640                .as_bytes()
641                .get(abs_idx - 1)
642                .is_some_and(|&b| b.is_ascii_alphanumeric() || b == b'_');
643
644        let after_idx = abs_idx + word_len;
645        let after_ok = after_idx >= text.len()
646            || !text
647                .as_bytes()
648                .get(after_idx)
649                .is_some_and(|&b| b.is_ascii_alphanumeric() || b == b'_');
650
651        if before_ok && after_ok {
652            return Some(abs_idx);
653        }
654        pos = abs_idx + 1;
655    }
656    None
657}
658
659/// Parse INSERTED.col1, INSERTED.col2, ... and extract column names as "col1, col2, ..."
660fn parse_inserted_columns(output_str: &str) -> String {
661    let mut cols = Vec::new();
662
663    for part in output_str.split(',') {
664        let trimmed = part.trim();
665
666        if let Some(after_inserted) = trimmed
667            .strip_prefix("INSERTED.")
668            .or_else(|| trimmed.strip_prefix("inserted."))
669            .or_else(|| trimmed.strip_prefix("INSERTED"))
670            .or_else(|| trimmed.strip_prefix("inserted"))
671        {
672            let col_name = after_inserted.trim().to_string();
673            if !col_name.is_empty() {
674                cols.push(col_name);
675            }
676        }
677    }
678
679    cols.join(", ")
680}
681
682/// Strip the `INTO :N, :N, ...` suffix from an Oracle `RETURNING ... INTO` clause.
683fn strip_returning_into(sql: &str) -> String {
684    let upper = sql.to_uppercase();
685    if let Some(ret_pos) = upper.rfind("RETURNING") {
686        let after_returning = &upper[ret_pos + "RETURNING".len()..];
687        if let Some(into_offset) = after_returning.find("INTO") {
688            let into_pos = ret_pos + "RETURNING".len() + into_offset;
689            let trimmed = sql[..into_pos].trim_end();
690            return trimmed.to_string();
691        }
692    }
693    sql.to_string()
694}
695
696#[cfg(test)]
697mod tests {
698    use super::*;
699    use crate::errors::ErrorCode;
700
701    fn parse(sql: &str) -> Result<Query, ScytheError> {
702        parse_query(sql)
703    }
704
705    #[test]
706    fn test_basic_parse() {
707        let input = "-- @name GetUsers\n-- @returns :many\nSELECT * FROM users;";
708        let q = parse(input).unwrap();
709        assert_eq!(q.name, "GetUsers");
710        assert_eq!(q.command, QueryCommand::Many);
711        assert!(q.sql.contains("SELECT"));
712    }
713
714    #[test]
715    fn test_all_command_types() {
716        let cases = vec![
717            (":one", QueryCommand::One),
718            (":many", QueryCommand::Many),
719            (":exec", QueryCommand::Exec),
720            (":exec_result", QueryCommand::ExecResult),
721            (":exec_rows", QueryCommand::ExecRows),
722        ];
723        for (tag, expected) in cases {
724            let input = format!("-- @name Q\n-- @returns {}\nSELECT 1", tag);
725            let q = parse(&input).unwrap();
726            assert_eq!(q.command, expected, "failed for {}", tag);
727        }
728    }
729
730    #[test]
731    fn test_case_insensitive_keywords() {
732        let input = "-- @Name GetUsers\n-- @RETURNS :many\nSELECT 1";
733        let q = parse(input).unwrap();
734        assert_eq!(q.name, "GetUsers");
735        assert_eq!(q.command, QueryCommand::Many);
736    }
737
738    #[test]
739    fn test_missing_name_errors() {
740        let input = "-- @returns :many\nSELECT 1";
741        let err = parse(input).unwrap_err();
742        assert_eq!(err.code, ErrorCode::MissingAnnotation);
743        assert!(err.message.contains("name"));
744    }
745
746    #[test]
747    fn test_missing_returns_errors() {
748        let input = "-- @name Foo\nSELECT 1";
749        let err = parse(input).unwrap_err();
750        assert_eq!(err.code, ErrorCode::MissingAnnotation);
751        assert!(err.message.contains("returns"));
752    }
753
754    #[test]
755    fn test_invalid_returns_value() {
756        let input = "-- @name Foo\n-- @returns :invalid\nSELECT 1";
757        let err = parse(input).unwrap_err();
758        assert_eq!(err.code, ErrorCode::InvalidAnnotation);
759    }
760
761    #[test]
762    fn test_empty_name_value() {
763        let input = "-- @name\n-- @returns :one\nSELECT 1";
764        let q = parse(input).unwrap();
765        assert_eq!(q.name, "");
766    }
767
768    #[test]
769    fn test_param_annotation() {
770        let input = "-- @name Foo\n-- @returns :one\n-- @param id: the user ID\nSELECT 1";
771        let q = parse(input).unwrap();
772        assert_eq!(q.annotations.param_docs.len(), 1);
773        assert_eq!(q.annotations.param_docs[0].name, "id");
774        assert_eq!(q.annotations.param_docs[0].description, "the user ID");
775    }
776
777    #[test]
778    fn test_param_no_description() {
779        let input = "-- @name Foo\n-- @returns :one\n-- @param id\nSELECT 1";
780        let q = parse(input).unwrap();
781        assert_eq!(q.annotations.param_docs.len(), 1);
782        assert_eq!(q.annotations.param_docs[0].name, "id");
783        assert_eq!(q.annotations.param_docs[0].description, "");
784    }
785
786    #[test]
787    fn test_nullable_annotation() {
788        let input = "-- @name Foo\n-- @returns :one\n-- @nullable col1, col2\nSELECT 1";
789        let q = parse(input).unwrap();
790        assert_eq!(q.annotations.nullable_overrides, vec!["col1", "col2"]);
791    }
792
793    #[test]
794    fn test_nonnull_annotation() {
795        let input = "-- @name Foo\n-- @returns :one\n-- @nonnull col1\nSELECT 1";
796        let q = parse(input).unwrap();
797        assert_eq!(q.annotations.nonnull_overrides, vec!["col1"]);
798    }
799
800    #[test]
801    fn test_json_annotation() {
802        let input = "-- @name Foo\n-- @returns :one\n-- @json data = EventData\nSELECT 1";
803        let q = parse(input).unwrap();
804        assert_eq!(q.annotations.json_mappings.len(), 1);
805        assert_eq!(q.annotations.json_mappings[0].column, "data");
806        assert_eq!(q.annotations.json_mappings[0].rust_type, "EventData");
807    }
808
809    #[test]
810    fn test_custom_annotations_captured() {
811        let input = "-- @name GetUser
812-- @returns :one
813-- @http GET /users/{id}
814-- @http_auth bearer:jwt
815-- @http_status 200,404
816SELECT id FROM users WHERE id = $1";
817        let q = parse(input).unwrap();
818        assert_eq!(q.annotations.custom.len(), 3);
819        assert_eq!(q.annotations.custom[0].name, "http");
820        assert_eq!(q.annotations.custom[0].value, "GET /users/{id}");
821        assert_eq!(q.annotations.custom[0].line, 3);
822        assert_eq!(q.annotations.custom[1].name, "http_auth");
823        assert_eq!(q.annotations.custom[1].value, "bearer:jwt");
824        assert_eq!(q.annotations.custom[1].line, 4);
825        assert_eq!(q.annotations.custom[2].name, "http_status");
826        assert_eq!(q.annotations.custom[2].value, "200,404");
827        assert_eq!(q.annotations.custom[2].line, 5);
828    }
829
830    #[test]
831    fn test_custom_annotation_without_value() {
832        let input = "-- @name GetUser
833-- @returns :one
834-- @http_internal
835SELECT 1";
836        let q = parse(input).unwrap();
837        assert_eq!(q.annotations.custom.len(), 1);
838        assert_eq!(q.annotations.custom[0].name, "http_internal");
839        assert_eq!(q.annotations.custom[0].value, "");
840    }
841
842    #[cfg(feature = "serde")]
843    #[test]
844    fn test_custom_annotation_serde_round_trip() {
845        let original = CustomAnnotation {
846            name: "http".to_string(),
847            value: "GET /users/{id}".to_string(),
848            line: 7,
849        };
850        let json = serde_json::to_string(&original).unwrap();
851        let back: CustomAnnotation = serde_json::from_str(&json).unwrap();
852        assert_eq!(back, original);
853    }
854
855    #[test]
856    fn test_custom_annotation_name_lowercased() {
857        let input = "-- @name GetUser
858-- @returns :one
859-- @HTTP_Auth Bearer
860SELECT 1";
861        let q = parse(input).unwrap();
862        assert_eq!(q.annotations.custom.len(), 1);
863        assert_eq!(q.annotations.custom[0].name, "http_auth");
864        assert_eq!(q.annotations.custom[0].value, "Bearer");
865    }
866
867    // ---- @param positional form ----
868
869    #[test]
870    fn test_positional_param_basic() {
871        let input = "-- @name Foo\n-- @returns :one\n-- @param $1 user_id\nSELECT 1";
872        let q = parse(input).unwrap();
873        assert_eq!(q.annotations.positional_param_docs.len(), 1);
874        assert_eq!(q.annotations.positional_param_docs[0].position, 1);
875        assert_eq!(q.annotations.positional_param_docs[0].name, "user_id");
876        assert_eq!(q.annotations.positional_param_docs[0].description, "");
877        assert_eq!(q.annotations.param_docs.len(), 0);
878    }
879
880    #[test]
881    fn test_positional_param_with_description() {
882        let input = "-- @name Foo\n-- @returns :one\n-- @param $4 bucket: time bucket as text\nSELECT 1";
883        let q = parse(input).unwrap();
884        assert_eq!(q.annotations.positional_param_docs.len(), 1);
885        assert_eq!(q.annotations.positional_param_docs[0].position, 4);
886        assert_eq!(q.annotations.positional_param_docs[0].name, "bucket");
887        assert_eq!(
888            q.annotations.positional_param_docs[0].description,
889            "time bucket as text"
890        );
891    }
892
893    #[test]
894    fn test_positional_param_does_not_affect_docs_only_param() {
895        let input = "-- @name Foo\n-- @returns :one\n-- @param id: the user ID\n-- @param $2 name\nSELECT 1";
896        let q = parse(input).unwrap();
897        assert_eq!(q.annotations.param_docs.len(), 1);
898        assert_eq!(q.annotations.param_docs[0].name, "id");
899        assert_eq!(q.annotations.positional_param_docs.len(), 1);
900        assert_eq!(q.annotations.positional_param_docs[0].position, 2);
901        assert_eq!(q.annotations.positional_param_docs[0].name, "name");
902    }
903
904    #[test]
905    fn test_positional_param_multiple() {
906        let input =
907            "-- @name Foo\n-- @returns :one\n-- @param $1 start_date: start\n-- @param $2 end_date: end\nSELECT 1";
908        let q = parse(input).unwrap();
909        assert_eq!(q.annotations.positional_param_docs.len(), 2);
910        assert_eq!(q.annotations.positional_param_docs[0].position, 1);
911        assert_eq!(q.annotations.positional_param_docs[0].name, "start_date");
912        assert_eq!(q.annotations.positional_param_docs[1].position, 2);
913        assert_eq!(q.annotations.positional_param_docs[1].name, "end_date");
914    }
915
916    #[test]
917    fn test_deprecated_annotation() {
918        let input = "-- @name Foo\n-- @returns :one\n-- @deprecated Use V2\nSELECT 1";
919        let q = parse(input).unwrap();
920        assert_eq!(q.annotations.deprecated, Some("Use V2".to_string()));
921    }
922
923    #[test]
924    fn test_sql_syntax_error() {
925        let input = "-- @name Foo\n-- @returns :one\nSELCT * FROM users";
926        let err = parse(input).unwrap_err();
927        assert_eq!(err.code, ErrorCode::SyntaxError);
928    }
929
930    #[test]
931    fn test_trailing_semicolon() {
932        let input = "-- @name Foo\n-- @returns :one\nSELECT 1;";
933        let q = parse(input).unwrap();
934        assert_eq!(q.name, "Foo");
935    }
936
937    #[test]
938    fn test_multiple_statements_error() {
939        let input = "-- @name Foo\n-- @returns :one\nSELECT 1; SELECT 2;";
940        let err = parse(input).unwrap_err();
941        assert_eq!(err.code, ErrorCode::SyntaxError);
942    }
943
944    #[test]
945    fn test_sql_preserved_without_annotations() {
946        let input = "-- @name Foo\n-- @returns :one\nSELECT id, name FROM users WHERE id = $1";
947        let q = parse(input).unwrap();
948        assert_eq!(q.sql, "SELECT id, name FROM users WHERE id = $1");
949    }
950
951    #[test]
952    fn test_returns_without_colon_prefix() {
953        let input = "-- @name Foo\n-- @returns many\nSELECT 1";
954        let q = parse(input).unwrap();
955        assert_eq!(q.command, QueryCommand::Many);
956    }
957
958    #[test]
959    fn test_batch_command() {
960        let input = "-- @name Foo\n-- @returns :batch\nSELECT 1";
961        let q = parse(input).unwrap();
962        assert_eq!(q.command, QueryCommand::Batch);
963    }
964
965    #[test]
966    fn test_grouped_command_with_group_by() {
967        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";
968        let q = parse(input).unwrap();
969        assert_eq!(q.command, QueryCommand::Grouped);
970        assert_eq!(q.annotations.group_by, Some("users.id".to_string()));
971    }
972
973    #[test]
974    fn test_grouped_command_without_group_by_errors() {
975        let input = "-- @name Foo\n-- @returns :grouped\nSELECT 1";
976        let err = parse(input).unwrap_err();
977        assert_eq!(err.code, ErrorCode::InvalidAnnotation);
978        assert!(err.message.contains("@group_by"));
979    }
980
981    #[test]
982    fn test_group_by_without_grouped_is_ignored() {
983        let input = "-- @name Foo\n-- @returns :many\n-- @group_by users.id\nSELECT 1";
984        let q = parse(input).unwrap();
985        assert_eq!(q.command, QueryCommand::Many);
986        assert_eq!(q.annotations.group_by, Some("users.id".to_string()));
987    }
988
989    #[test]
990    fn test_preprocess_postgres_strips_partial_index_where() {
991        let sql = "INSERT INTO billing_events (project_id, stripe_event_id) \
992                   VALUES ($1, $2) \
993                   ON CONFLICT (stripe_event_id) WHERE stripe_event_id IS NOT NULL DO NOTHING";
994        let cleaned = preprocess_postgres_sql(sql);
995        assert!(
996            !cleaned.to_uppercase().contains("WHERE STRIPE_EVENT_ID IS NOT NULL"),
997            "WHERE clause must be stripped between ON CONFLICT cols and DO; got: {cleaned}"
998        );
999        assert!(
1000            cleaned
1001                .to_uppercase()
1002                .contains("ON CONFLICT (STRIPE_EVENT_ID) DO NOTHING")
1003        );
1004        sqlparser::parser::Parser::parse_sql(&sqlparser::dialect::PostgreSqlDialect {}, &cleaned)
1005            .expect("cleaned SQL should parse");
1006    }
1007
1008    #[test]
1009    fn test_preprocess_postgres_no_op_when_no_partial_clause() {
1010        let sql = "INSERT INTO t (a) VALUES ($1) ON CONFLICT (a) DO UPDATE SET a = EXCLUDED.a";
1011        assert_eq!(preprocess_postgres_sql(sql), sql);
1012    }
1013
1014    #[test]
1015    fn test_preprocess_postgres_leaves_on_conflict_on_constraint_alone() {
1016        let sql = "INSERT INTO t (a) VALUES ($1) ON CONFLICT ON CONSTRAINT t_a_uidx DO NOTHING";
1017        assert_eq!(preprocess_postgres_sql(sql), sql);
1018    }
1019
1020    #[test]
1021    fn test_preprocess_postgres_handles_compound_index_cols() {
1022        let sql = "INSERT INTO t (a, b) VALUES ($1, $2) \
1023                   ON CONFLICT (a, b) WHERE a IS NOT NULL AND b > 0 DO UPDATE SET b = EXCLUDED.b";
1024        let cleaned = preprocess_postgres_sql(sql);
1025        assert!(cleaned.to_uppercase().contains("ON CONFLICT (A, B) DO UPDATE"));
1026        assert!(!cleaned.to_uppercase().contains("WHERE A IS NOT NULL"));
1027    }
1028
1029    #[test]
1030    fn test_preprocess_postgres_preserves_unrelated_where() {
1031        let sql = "DELETE FROM t WHERE id = $1";
1032        assert_eq!(preprocess_postgres_sql(sql), sql);
1033    }
1034
1035    #[test]
1036    fn test_preprocess_postgres_ignores_text_inside_line_comments() {
1037        let sql = "-- inline doc: `ON CONFLICT (col) WHERE …` is the partial form\n\
1038                   INSERT INTO t (a) VALUES ($1) \
1039                   ON CONFLICT (a) WHERE a IS NOT NULL DO NOTHING";
1040        let cleaned = preprocess_postgres_sql(sql);
1041        assert!(
1042            cleaned.contains("-- inline doc"),
1043            "comment must survive the pass; got: {cleaned}"
1044        );
1045        assert!(cleaned.contains("ON CONFLICT (a) DO NOTHING"));
1046    }
1047
1048    #[test]
1049    fn test_preprocess_postgres_ignores_text_inside_string_literals() {
1050        let sql = "SELECT 'ON CONFLICT (a) WHERE a IS NOT NULL DO NOTHING' AS s";
1051        assert_eq!(preprocess_postgres_sql(sql), sql);
1052    }
1053
1054    #[test]
1055    fn test_preprocess_oracle_colon_placeholders() {
1056        assert_eq!(
1057            preprocess_oracle_sql("SELECT * FROM users WHERE id = :1"),
1058            "SELECT * FROM users WHERE id = ?"
1059        );
1060        assert_eq!(
1061            preprocess_oracle_sql("INSERT INTO users (name, email) VALUES (:1, :2)"),
1062            "INSERT INTO users (name, email) VALUES (?, ?)"
1063        );
1064    }
1065
1066    #[test]
1067    fn test_preprocess_oracle_preserves_string_literals() {
1068        assert_eq!(
1069            preprocess_oracle_sql("SELECT * FROM users WHERE name = ':1' AND id = :1"),
1070            "SELECT * FROM users WHERE name = ':1' AND id = ?"
1071        );
1072    }
1073
1074    #[test]
1075    fn test_preprocess_oracle_strips_returning_into() {
1076        assert_eq!(
1077            preprocess_oracle_sql("INSERT INTO users (name) VALUES (:1) RETURNING id, name INTO :2, :3"),
1078            "INSERT INTO users (name) VALUES (?) RETURNING id, name"
1079        );
1080    }
1081
1082    #[test]
1083    fn test_preprocess_oracle_full_insert_returning_into() {
1084        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";
1085        let result = preprocess_oracle_sql(sql);
1086        assert_eq!(
1087            result,
1088            "INSERT INTO users (name, email, active) VALUES (?, ?, ?) RETURNING id, name, email, active, created_at"
1089        );
1090    }
1091
1092    #[test]
1093    fn test_preprocess_oracle_no_returning_into_unchanged() {
1094        assert_eq!(
1095            preprocess_oracle_sql("DELETE FROM users WHERE id = :1"),
1096            "DELETE FROM users WHERE id = ?"
1097        );
1098    }
1099
1100    #[test]
1101    fn test_preprocess_mssql_single_placeholder() {
1102        assert_eq!(
1103            preprocess_mssql_sql("SELECT * FROM users WHERE id = @p1"),
1104            "SELECT * FROM users WHERE id = ?"
1105        );
1106    }
1107
1108    #[test]
1109    fn test_preprocess_mssql_multiple_placeholders() {
1110        assert_eq!(
1111            preprocess_mssql_sql("INSERT INTO users (name, email) VALUES (@p1, @p2)"),
1112            "INSERT INTO users (name, email) VALUES (?, ?)"
1113        );
1114    }
1115
1116    #[test]
1117    fn test_preprocess_mssql_preserves_string_literals() {
1118        assert_eq!(
1119            preprocess_mssql_sql("SELECT * FROM users WHERE name = '@p1' AND id = @p1"),
1120            "SELECT * FROM users WHERE name = '@p1' AND id = ?"
1121        );
1122    }
1123
1124    #[test]
1125    fn test_preprocess_mssql_case_insensitive_p() {
1126        assert_eq!(
1127            preprocess_mssql_sql("SELECT * FROM users WHERE id = @P1"),
1128            "SELECT * FROM users WHERE id = ?"
1129        );
1130    }
1131
1132    #[test]
1133    fn test_preprocess_mssql_non_placeholder_at_variable_unchanged() {
1134        assert_eq!(preprocess_mssql_sql("SELECT @myvar"), "SELECT @myvar");
1135    }
1136
1137    #[test]
1138    fn test_preprocess_mssql_multi_digit_placeholder() {
1139        assert_eq!(preprocess_mssql_sql("SELECT @p10, @p2"), "SELECT ?, ?");
1140    }
1141
1142    #[test]
1143    fn test_preprocess_mssql_output_inserted_simple() {
1144        let sql = "INSERT INTO users (id, name) OUTPUT INSERTED.id, INSERTED.name VALUES (@p1, @p2)";
1145        let result = preprocess_mssql_sql(sql);
1146        assert!(result.contains("RETURNING id, name"), "got: {}", result);
1147        assert!(result.contains("VALUES (?, ?)"), "got: {}", result);
1148        assert!(!result.contains("OUTPUT"), "got: {}", result);
1149    }
1150
1151    #[test]
1152    fn test_preprocess_mssql_output_inserted_full_example() {
1153        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)";
1154        let result = preprocess_mssql_sql(sql);
1155        assert!(
1156            result.contains("RETURNING id, name, email, active, created_at"),
1157            "got: {}",
1158            result
1159        );
1160        assert!(result.contains("VALUES (?, ?, ?, ?)"), "got: {}", result);
1161    }
1162
1163    #[test]
1164    fn test_preprocess_mssql_output_case_insensitive() {
1165        let sql = "INSERT INTO users (id) output inserted.id values (@p1)";
1166        let result = preprocess_mssql_sql(sql);
1167        assert!(result.contains("RETURNING id"), "got: {}", result);
1168        assert!(
1169            result.contains("values (?)") || result.contains("VALUES (?)"),
1170            "got: {}",
1171            result
1172        );
1173    }
1174
1175    #[test]
1176    fn test_preprocess_mssql_no_output_unchanged() {
1177        let sql = "INSERT INTO users (id, name) VALUES (@p1, @p2)";
1178        let result = preprocess_mssql_sql(sql);
1179        assert_eq!(result, "INSERT INTO users (id, name) VALUES (?, ?)");
1180    }
1181
1182    #[test]
1183    fn test_preprocess_mssql_output_with_string_literal() {
1184        let sql = "INSERT INTO users (id, name) OUTPUT INSERTED.id, INSERTED.name VALUES (@p1, '@p2')";
1185        let result = preprocess_mssql_sql(sql);
1186        assert!(result.contains("RETURNING id, name"), "got: {}", result);
1187        assert!(result.contains("(?, '@p2')"), "got: {}", result);
1188    }
1189
1190    #[test]
1191    fn test_preprocess_mssql_output_with_whitespace() {
1192        let sql = "INSERT INTO users (id, name)\nOUTPUT INSERTED.id,\n  INSERTED.name\nVALUES (@p1, @p2)";
1193        let result = preprocess_mssql_sql(sql);
1194        assert!(result.contains("RETURNING id, name"), "got: {}", result);
1195        assert!(result.contains("VALUES (?, ?)"), "got: {}", result);
1196    }
1197}
scythe_core/parser/mod.rs

scythe_core/parser/
mod.rs
