rullst-orm 12.1.1

//! SQL-aware conversion of portable `?` bind markers to PostgreSQL `$n` markers.

/// Converts portable `?` bind markers to PostgreSQL `$1`, `$2`, ... markers.
///
/// Quoted strings and identifiers, line and nested block comments, PostgreSQL
/// dollar-quoted strings, existing `$n` parameters, and the unambiguous JSON
/// operators `?|`, `?&`, and `@?` are recognized. This keeps SQL syntax and
/// comments intact while converting bind markers generated by Rullst.
#[doc(hidden)]
pub fn replace_placeholders(sql: &str) -> String {
    let scan = scan_sql(sql.as_bytes());
    if scan.placeholders.is_empty() {
        return sql.to_string();
    }

    let extra_capacity = scan.placeholders.len().saturating_mul(3);
    let mut replaced = String::with_capacity(sql.len().saturating_add(extra_capacity));
    let mut last_index = 0;
    let mut parameter = scan.highest_parameter.saturating_add(1).max(1);

    for index in scan.placeholders {
        replaced.push_str(&sql[last_index..index]);
        replaced.push('$');
        replaced.push_str(&parameter.to_string());
        parameter = parameter.saturating_add(1);
        last_index = index + 1;
    }
    replaced.push_str(&sql[last_index..]);
    replaced
}

struct SqlScan {
    placeholders: Vec<usize>,
    highest_parameter: usize,
}

#[derive(Clone, Copy)]
enum LexState {
    Code,
    SingleQuoted,
    DoubleQuoted,
    BacktickQuoted,
    LineComment,
    BlockComment(usize),
    DollarQuoted {
        delimiter_start: usize,
        delimiter_len: usize,
    },
}

fn scan_sql(sql: &[u8]) -> SqlScan {
    let mut placeholders = Vec::new();
    let mut highest_parameter = 0;
    let mut state = LexState::Code;
    let mut index = 0;

    while index < sql.len() {
        match state {
            LexState::Code => {
                if starts_with(sql, index, b"--") {
                    state = LexState::LineComment;
                    index += 2;
                } else if starts_with(sql, index, b"/*") {
                    state = LexState::BlockComment(1);
                    index += 2;
                } else {
                    match sql[index] {
                        b'\'' => {
                            state = LexState::SingleQuoted;
                            index += 1;
                        }
                        b'"' => {
                            state = LexState::DoubleQuoted;
                            index += 1;
                        }
                        b'`' => {
                            state = LexState::BacktickQuoted;
                            index += 1;
                        }
                        b'$' => {
                            if let Some((value, length)) = numbered_parameter(sql, index) {
                                highest_parameter = highest_parameter.max(value);
                                index += length;
                            } else if let Some(delimiter_len) = dollar_delimiter_len(sql, index) {
                                state = LexState::DollarQuoted {
                                    delimiter_start: index,
                                    delimiter_len,
                                };
                                index += delimiter_len;
                            } else {
                                index += 1;
                            }
                        }
                        b'?' => {
                            if !is_question_operator(sql, index) {
                                placeholders.push(index);
                            }
                            index += 1;
                        }
                        _ => index += 1,
                    }
                }
            }
            LexState::SingleQuoted => {
                if sql[index] == b'\\' {
                    index = index.saturating_add(2).min(sql.len());
                } else if sql[index] == b'\'' {
                    if sql.get(index + 1) == Some(&b'\'') {
                        index += 2;
                    } else {
                        state = LexState::Code;
                        index += 1;
                    }
                } else {
                    index += 1;
                }
            }
            LexState::DoubleQuoted => {
                if sql[index] == b'"' {
                    if sql.get(index + 1) == Some(&b'"') {
                        index += 2;
                    } else {
                        state = LexState::Code;
                        index += 1;
                    }
                } else {
                    index += 1;
                }
            }
            LexState::BacktickQuoted => {
                if sql[index] == b'`' {
                    if sql.get(index + 1) == Some(&b'`') {
                        index += 2;
                    } else {
                        state = LexState::Code;
                        index += 1;
                    }
                } else {
                    index += 1;
                }
            }
            LexState::LineComment => {
                if sql[index] == b'\n' {
                    state = LexState::Code;
                }
                index += 1;
            }
            LexState::BlockComment(depth) => {
                if starts_with(sql, index, b"/*") {
                    state = LexState::BlockComment(depth.saturating_add(1));
                    index += 2;
                } else if starts_with(sql, index, b"*/") {
                    if depth == 1 {
                        state = LexState::Code;
                    } else {
                        state = LexState::BlockComment(depth - 1);
                    }
                    index += 2;
                } else {
                    index += 1;
                }
            }
            LexState::DollarQuoted {
                delimiter_start,
                delimiter_len,
            } => {
                let delimiter_end = delimiter_start.saturating_add(delimiter_len);
                let delimiter = &sql[delimiter_start..delimiter_end];
                if starts_with(sql, index, delimiter) {
                    state = LexState::Code;
                    index += delimiter_len;
                } else {
                    index += 1;
                }
            }
        }
    }

    SqlScan {
        placeholders,
        highest_parameter,
    }
}

fn starts_with(sql: &[u8], index: usize, token: &[u8]) -> bool {
    sql.get(index..)
        .is_some_and(|remaining| remaining.starts_with(token))
}

fn numbered_parameter(sql: &[u8], index: usize) -> Option<(usize, usize)> {
    let first_digit = *sql.get(index + 1)?;
    if !first_digit.is_ascii_digit() {
        return None;
    }

    let mut cursor = index + 1;
    let mut value = 0_usize;
    while let Some(byte) = sql.get(cursor)
        && byte.is_ascii_digit()
    {
        value = value
            .saturating_mul(10)
            .saturating_add(usize::from(*byte - b'0'));
        cursor += 1;
    }
    Some((value, cursor - index))
}

fn dollar_delimiter_len(sql: &[u8], index: usize) -> Option<usize> {
    let first = *sql.get(index + 1)?;
    if first == b'$' {
        return Some(2);
    }
    if !first.is_ascii_alphabetic() && first != b'_' {
        return None;
    }

    let mut cursor = index + 2;
    while let Some(byte) = sql.get(cursor) {
        if *byte == b'$' {
            return Some(cursor - index + 1);
        }
        if !byte.is_ascii_alphanumeric() && *byte != b'_' {
            return None;
        }
        cursor += 1;
    }
    None
}

fn is_question_operator(sql: &[u8], index: usize) -> bool {
    if matches!(sql.get(index + 1), Some(b'|' | b'&'))
        || sql.get(index.wrapping_sub(1)) == Some(&b'@')
    {
        return true;
    }

    let previous = previous_non_whitespace(sql, index);
    let next = next_non_whitespace(sql, index + 1);
    matches!(next, Some(b'\'' | b'"' | b'$'))
        && previous.is_some_and(|byte| {
            byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b')' | b']' | b'\'' | b'"')
        })
}

fn previous_non_whitespace(sql: &[u8], before: usize) -> Option<u8> {
    let mut cursor = before;
    while cursor > 0 {
        cursor -= 1;
        if !sql[cursor].is_ascii_whitespace() {
            return Some(sql[cursor]);
        }
    }
    None
}

fn next_non_whitespace(sql: &[u8], after: usize) -> Option<u8> {
    sql.get(after..)?
        .iter()
        .copied()
        .find(|byte| !byte.is_ascii_whitespace())
}

#[cfg(test)]
mod tests {
    use super::replace_placeholders;

    #[test]
    fn preserves_quoted_text_comments_and_dollar_quotes() {
        let sql = "SELECT '?', \"?\", `?`, $$?$$, $body$?$body$ /* ? /* ? */ */ FROM t -- ?\nWHERE id = ? AND name = ?";
        assert_eq!(
            replace_placeholders(sql),
            "SELECT '?', \"?\", `?`, $$?$$, $body$?$body$ /* ? /* ? */ */ FROM t -- ?\nWHERE id = $1 AND name = $2"
        );
    }

    #[test]
    fn preserves_json_operators_and_continues_existing_parameters() {
        let sql = "SELECT payload ? 'key', payload ?| array['a'], payload ?& array['b'], payload @? '$.a' FROM events WHERE tenant_id = $4 AND id = ?";
        assert_eq!(
            replace_placeholders(sql),
            "SELECT payload ? 'key', payload ?| array['a'], payload ?& array['b'], payload @? '$.a' FROM events WHERE tenant_id = $4 AND id = $5"
        );
    }
}