rullst-orm 12.3.0

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

use crate::{Error, RullstValue};

/// 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. Strings follow PostgreSQL's
/// default `standard_conforming_strings = on`: a backslash escapes the next
/// character only inside `E'...'` escape strings. 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
}

/// Rewrites one embedded subquery fragment to portable `?` bind markers.
///
/// A typed subquery renders PostgreSQL `$n` markers when the active driver is
/// PostgreSQL, but the containing builder must number every marker of the
/// final statement exactly once, in textual order. Each `$n` is therefore
/// replaced by `?` and its binding is emitted at that textual position, so
/// repeated or out-of-order markers from custom builders keep their meaning.
/// Fragments that already use `?` markers are returned unchanged.
///
/// A fragment that mixes `?` and `$n` markers, references a missing binding,
/// or leaves a supplied binding unreferenced is rejected instead of guessed.
#[doc(hidden)]
pub fn portable_subquery(
    sql: &str,
    bindings: Vec<RullstValue>,
) -> Result<(String, Vec<RullstValue>), Error> {
    let scan = scan_sql(sql.as_bytes());
    if scan.numbered.is_empty() {
        return Ok((sql.to_string(), bindings));
    }
    if !scan.placeholders.is_empty() {
        return Err(Error::Validation(
            "subquery mixes portable `?` and numbered `$n` bind markers".to_string(),
        ));
    }

    let mut portable = String::with_capacity(sql.len());
    let mut ordered = Vec::with_capacity(scan.numbered.len());
    let mut referenced = vec![false; bindings.len()];
    let mut last_index = 0;
    for marker in &scan.numbered {
        let binding = marker
            .value
            .checked_sub(1)
            .and_then(|index| bindings.get(index).map(|value| (index, value)));
        let Some((index, value)) = binding else {
            return Err(Error::Validation(format!(
                "subquery bind marker ${} has no matching binding ({} supplied)",
                marker.value,
                bindings.len()
            )));
        };
        referenced[index] = true;
        ordered.push(value.clone());
        portable.push_str(&sql[last_index..marker.start]);
        portable.push('?');
        last_index = marker.start + marker.length;
    }
    if referenced.iter().any(|used| !used) {
        return Err(Error::Validation(
            "subquery supplies a binding that its SQL never references".to_string(),
        ));
    }
    portable.push_str(&sql[last_index..]);
    Ok((portable, ordered))
}

struct NumberedParameter {
    start: usize,
    length: usize,
    value: usize,
}

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

#[derive(Clone, Copy)]
enum LexState {
    Code,
    /// `backslash_escapes` is set for PostgreSQL `E'...'` escape strings.
    SingleQuoted {
        backslash_escapes: bool,
    },
    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 numbered = 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 {
                                backslash_escapes: is_escape_string_prefix(sql, index),
                            };
                            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);
                                numbered.push(NumberedParameter {
                                    start: index,
                                    length,
                                    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 { backslash_escapes } => {
                if backslash_escapes && 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,
        numbered,
        highest_parameter,
    }
}

/// Whether the quote at `quote` opens a PostgreSQL escape string: it directly
/// follows an `E`/`e` that is not the tail of a longer identifier.
fn is_escape_string_prefix(sql: &[u8], quote: usize) -> bool {
    let Some(prefix) = quote.checked_sub(1) else {
        return false;
    };
    if !matches!(sql[prefix], b'e' | b'E') {
        return false;
    }
    prefix
        .checked_sub(1)
        .and_then(|before| sql.get(before))
        .is_none_or(|byte| {
            !(byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'$') || *byte >= 0x80)
        })
}

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::{portable_subquery, replace_placeholders};
    use crate::RullstValue;

    fn text(value: &str) -> RullstValue {
        RullstValue::String(value.to_string())
    }

    fn described(values: &[RullstValue]) -> Vec<String> {
        values.iter().map(|value| format!("{value:?}")).collect()
    }

    #[test]
    fn portable_subquery_restores_textual_marker_order() {
        let rendered = replace_placeholders(
            "SELECT * FROM comments WHERE (tenant_id = ?) AND (status = ?) AND body = '$9'",
        );
        let (sql, bindings) =
            portable_subquery(&rendered, vec![text("acme"), text("open")]).expect("portable");
        assert_eq!(
            sql,
            "SELECT * FROM comments WHERE (tenant_id = ?) AND (status = ?) AND body = '$9'"
        );
        assert_eq!(
            described(&bindings),
            described(&[text("acme"), text("open")])
        );

        let (sql, bindings) = portable_subquery(
            "SELECT 1 WHERE b = $2 AND a = $1 AND c = $2",
            vec![text("a"), text("b")],
        )
        .expect("custom numbering is reordered");
        assert_eq!(sql, "SELECT 1 WHERE b = ? AND a = ? AND c = ?");
        assert_eq!(
            described(&bindings),
            described(&[text("b"), text("a"), text("b")])
        );

        let (sql, bindings) =
            portable_subquery("SELECT 1 WHERE a = ?", vec![text("a")]).expect("already portable");
        assert_eq!(sql, "SELECT 1 WHERE a = ?");
        assert_eq!(described(&bindings), described(&[text("a")]));
    }

    #[test]
    fn portable_subquery_rejects_ambiguous_numbering() {
        for (sql, bindings) in [
            (
                "SELECT 1 WHERE a = ? AND b = $1",
                vec![text("a"), text("b")],
            ),
            ("SELECT 1 WHERE a = $2", vec![text("a")]),
            ("SELECT 1 WHERE a = $0", vec![text("a")]),
            ("SELECT 1 WHERE a = $1", vec![text("a"), text("unused")]),
        ] {
            assert!(matches!(
                portable_subquery(sql, bindings),
                Err(crate::Error::Validation(_))
            ));
        }
    }

    #[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 backslashes_escape_only_inside_escape_strings() {
        // Under standard_conforming_strings '\' is a complete literal.
        assert_eq!(
            replace_placeholders("name LIKE ? ESCAPE '\\' AND tenant_id = ?"),
            "name LIKE $1 ESCAPE '\\' AND tenant_id = $2"
        );
        assert_eq!(
            replace_placeholders("code = '\\' AND a = ? AND note = 'x\\y?'"),
            "code = '\\' AND a = $1 AND note = 'x\\y?'"
        );
        // E'...' strings honour backslash escapes, including an escaped quote.
        assert_eq!(
            replace_placeholders("SELECT E'it\\'s ?', e'\\\\' WHERE a = ?"),
            "SELECT E'it\\'s ?', e'\\\\' WHERE a = $1"
        );
        // An identifier ending in `e` does not start an escape string.
        assert_eq!(
            replace_placeholders("SELECT typee'\\' WHERE a = ?"),
            "SELECT typee'\\' WHERE a = $1"
        );
    }

    #[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"
        );
    }
}