skade-flight/src/odbc.rs
//! ODBC dialect shim — what a generic ODBC path (Power Query's
//! `Odbc.DataSource`, the Flight SQL ODBC driver) sends that DataFusion's
//! parser does not speak, rewritten BEFORE planning:
//!
//! - `SELECT TOP n …` → `SELECT … LIMIT n`
//! - `{d '2024-01-02'}` → `DATE '2024-01-02'`
//! - `{t '12:34:56'}` → `TIME '12:34:56'`
//! - `{ts '2024-01-02 12:34:56'}`→ `TIMESTAMP '2024-01-02 12:34:56'`
//! - `{fn NAME(args)}` → DataFusion's function, via the table below
//! (nested escapes handled, args rewritten)
//! - `{oj a LEFT OUTER JOIN b ON …}` → `a LEFT OUTER JOIN b ON …`
//! - `{escape '\'}` → `ESCAPE '\'`
//!
//! String literals (`'…'`, with `''`) and quoted identifiers (`"…"`) are
//! never touched. Anything not recognised is passed through untouched so the
//! planner reports it, not this shim. Idempotent: plain DataFusion SQL comes
//! out byte-identical.
/// Rewrite one statement. Plain SQL — no `{` and no `top` token — is
/// returned borrowed: zero allocation on the common path.
pub fn rewrite(sql: &str) -> std::borrow::Cow<'_, str> {
if !sql.contains('{') && !has_top(sql) {
return std::borrow::Cow::Borrowed(sql);
}
let braces = rewrite_braces(sql);
std::borrow::Cow::Owned(rewrite_top(&braces))
}
/// Case-insensitive whole-word `top` anywhere (cheap pre-check; the real
/// TOP parser decides). No allocation.
fn has_top(sql: &str) -> bool {
let b = sql.as_bytes();
let mut i = 0;
while i + 3 <= b.len() {
if b[i].eq_ignore_ascii_case(&b't')
&& b[i + 1].eq_ignore_ascii_case(&b'o')
&& b[i + 2].eq_ignore_ascii_case(&b'p')
&& word_boundary(b, i, i + 3)
{
return true;
}
i += 1;
}
false
}
// ── lexing helpers ──────────────────────────────────────────────────────────
/// Index just past a string literal starting at `i` (which is `'`).
fn skip_string(s: &[u8], mut i: usize) -> usize {
i += 1;
while i < s.len() {
if s[i] == b'\'' {
if i + 1 < s.len() && s[i + 1] == b'\'' {
i += 2;
continue;
}
return i + 1;
}
i += 1;
}
s.len()
}
/// Index just past a quoted identifier starting at `i` (which is `"`).
fn skip_quoted(s: &[u8], mut i: usize) -> usize {
i += 1;
while i < s.len() {
if s[i] == b'"' {
if i + 1 < s.len() && s[i + 1] == b'"' {
i += 2;
continue;
}
return i + 1;
}
i += 1;
}
s.len()
}
/// Find the `}` matching the `{` at `open`, respecting strings/quotes/nesting.
fn matching_brace(s: &[u8], open: usize) -> Option<usize> {
let mut depth = 0usize;
let mut i = open;
while i < s.len() {
match s[i] {
b'\'' => i = skip_string(s, i),
b'"' => i = skip_quoted(s, i),
b'{' => {
depth += 1;
i += 1;
}
b'}' => {
depth -= 1;
if depth == 0 {
return Some(i);
}
i += 1;
}
_ => i += 1,
}
}
None
}
/// Split a comma-separated argument list at top level (parens, strings).
fn split_args(s: &str) -> Vec<String> {
let b = s.as_bytes();
let mut out = Vec::new();
let mut depth = 0i32;
let mut start = 0usize;
let mut i = 0usize;
while i < b.len() {
match b[i] {
b'\'' => i = skip_string(b, i),
b'"' => i = skip_quoted(b, i),
b'(' | b'{' => {
depth += 1;
i += 1;
}
b')' | b'}' => {
depth -= 1;
i += 1;
}
b',' if depth == 0 => {
out.push(s[start..i].trim().to_string());
start = i + 1;
i += 1;
}
_ => i += 1,
}
}
let last = s[start..].trim();
if !last.is_empty() || !out.is_empty() {
out.push(last.to_string());
}
out
}
// ── {…} escapes ─────────────────────────────────────────────────────────────
fn rewrite_braces(sql: &str) -> String {
// Only ever inserts ASCII and copies input slices verbatim, so byte
// indices from the lexer are always on char boundaries.
let b = sql.as_bytes();
let mut out = String::with_capacity(sql.len());
let mut i = 0usize;
let mut plain_start = 0usize;
while i < b.len() {
match b[i] {
b'\'' => i = skip_string(b, i),
b'"' => i = skip_quoted(b, i),
b'{' => match matching_brace(b, i) {
Some(close) => {
out.push_str(&sql[plain_start..i]);
out.push_str(&rewrite_escape(&sql[i + 1..close]));
i = close + 1;
plain_start = i;
}
None => break,
},
_ => i += 1,
}
}
out.push_str(&sql[plain_start..]);
out
}
/// The inside of one `{…}` escape, already brace-stripped.
fn rewrite_escape(inner: &str) -> String {
let trimmed = inner.trim_start();
let lower = trimmed.to_ascii_lowercase();
let word_end = lower.find(|c: char| c.is_whitespace()).unwrap_or(lower.len());
let (kw, rest) = (&lower[..word_end], trimmed[word_end..].trim());
match kw {
"d" => format!("DATE {}", rewrite_braces(rest)),
"t" => format!("TIME {}", rewrite_braces(rest)),
"ts" => format!("TIMESTAMP {}", rewrite_braces(rest)),
"oj" => rewrite_braces(rest),
"escape" => format!("ESCAPE {}", rewrite_braces(rest)),
"fn" => rewrite_fn(rest),
// Unknown escape ({call …}, {interval …}, {guid …}): pass the body
// through and let the planner name it.
_ => rewrite_braces(trimmed),
}
}
/// `NAME(args)` from a `{fn …}` escape → DataFusion.
fn rewrite_fn(call: &str) -> String {
let Some(open) = call.find('(') else {
return rewrite_braces(call);
};
let name = call[..open].trim().to_ascii_uppercase();
let close = call.rfind(')').unwrap_or(call.len());
let args: Vec<String> = split_args(&call[open + 1..close]).into_iter().map(|a| rewrite_braces(&a)).collect();
let a = |i: usize| args.get(i).cloned().unwrap_or_default();
let plain = |f: &str| format!("{f}({})", args.join(", "));
let part = |unit: &str| format!("date_part('{unit}', {})", a(0));
match name.as_str() {
// string
"UCASE" => plain("upper"),
"LCASE" => plain("lower"),
"LENGTH" | "CHAR_LENGTH" | "CHARACTER_LENGTH" => plain("character_length"),
"OCTET_LENGTH" => plain("octet_length"),
"LOCATE" => match args.len() {
3 => format!("(CASE WHEN strpos(substr({1}, {2}), {0}) = 0 THEN 0 ELSE strpos(substr({1}, {2}), {0}) + {2} - 1 END)", a(0), a(1), a(2)),
_ => format!("strpos({}, {})", a(1), a(0)),
},
"POSITION" => plain("position"),
"LTRIM" => plain("ltrim"),
"RTRIM" => plain("rtrim"),
"TRIM" => plain("trim"),
"SUBSTRING" | "SUBSTR" => plain("substr"),
"LEFT" => plain("left"),
"RIGHT" => plain("right"),
"REPLACE" => plain("replace"),
"REPEAT" => plain("repeat"),
"CONCAT" => plain("concat"),
"ASCII" => plain("ascii"),
"CHAR" | "CHR" => plain("chr"),
"SPACE" => format!("repeat(' ', {})", a(0)),
"INSERT" => format!("concat(left({0}, {1} - 1), {3}, substr({0}, {1} + {2}))", a(0), a(1), a(2), a(3)),
"DIFFERENCE" | "SOUNDEX" => plain(&name.to_ascii_lowercase()),
// numeric
"MOD" => format!("(({}) % ({}))", a(0), a(1)),
"ABS" | "ACOS" | "ASIN" | "ATAN" | "ATAN2" | "CEILING" | "COS" | "COT" | "DEGREES" | "EXP" | "FLOOR"
| "LOG10" | "PI" | "POWER" | "RADIANS" | "ROUND" | "SIN" | "SQRT" | "TAN" => {
let f = match name.as_str() {
"CEILING" => "ceil",
other => &other.to_ascii_lowercase(),
};
plain(f)
}
"LOG" => plain("ln"),
"RAND" => "random()".to_string(),
"SIGN" => plain("signum"),
"TRUNCATE" => plain("trunc"),
// datetime
"NOW" | "CURRENT_TIMESTAMP" => "now()".to_string(),
"CURDATE" | "CURRENT_DATE" => "current_date()".to_string(),
"CURTIME" | "CURRENT_TIME" => "current_time()".to_string(),
"YEAR" => part("year"),
"QUARTER" => part("quarter"),
"MONTH" => part("month"),
"WEEK" => part("week"),
"DAYOFMONTH" => part("day"),
"DAYOFYEAR" => part("doy"),
"DAYOFWEEK" => format!("(date_part('dow', {}) + 1)", a(0)),
"HOUR" => part("hour"),
"MINUTE" => part("minute"),
"SECOND" => part("second"),
"EXTRACT" => plain("extract"),
"DAYNAME" => format!("to_char({}, '%A')", a(0)),
"MONTHNAME" => format!("to_char({}, '%B')", a(0)),
"TIMESTAMPDIFF" | "TIMESTAMPADD" => {
// {fn TIMESTAMPADD(SQL_TSI_DAY, 3, ts)} → ts + INTERVAL '3 day'
let unit = a(0).trim().to_ascii_uppercase();
let unit = unit.strip_prefix("SQL_TSI_").unwrap_or(&unit).to_ascii_lowercase();
let unit = if unit == "frac_second" { "microsecond".to_string() } else { unit };
if name == "TIMESTAMPDIFF" {
// TIMESTAMPDIFF(unit, a, b): whole units from a to b. Calendar
// units by parts, clock units by epoch seconds (integer division
// truncates toward zero, as ODBC specifies).
let (a0, b0) = (a(1), a(2));
return match unit.as_str() {
"year" => format!("(date_part('year', {b0}) - date_part('year', {a0}))"),
"quarter" => format!("(((date_part('year', {b0}) - date_part('year', {a0})) * 12 + (date_part('month', {b0}) - date_part('month', {a0}))) / 3)"),
"month" => format!("((date_part('year', {b0}) - date_part('year', {a0})) * 12 + (date_part('month', {b0}) - date_part('month', {a0})))"),
"week" => format!("((to_unixtime({b0}) - to_unixtime({a0})) / 604800)"),
"day" => format!("((to_unixtime({b0}) - to_unixtime({a0})) / 86400)"),
"hour" => format!("((to_unixtime({b0}) - to_unixtime({a0})) / 3600)"),
"minute" => format!("((to_unixtime({b0}) - to_unixtime({a0})) / 60)"),
"second" => format!("(to_unixtime({b0}) - to_unixtime({a0}))"),
// SQL_TSI_FRAC_SECOND is billionths of a second.
_ => format!("cast((date_part('epoch', {b0}) - date_part('epoch', {a0})) * 1000000000 as bigint)"),
};
}
if name == "TIMESTAMPADD" {
// DataFusion 55 cannot multiply Int64 by an interval, so a literal
// count becomes a literal interval and an expression goes through a
// string→interval cast.
let n = a(1);
let n = n.trim();
if n.chars().all(|c| c.is_ascii_digit() || c == '-') && !n.is_empty() {
format!("({} + INTERVAL '{n} {unit}')", a(2))
} else {
format!("({} + arrow_cast(concat(cast(({n}) as varchar), ' {unit}'), 'Interval(MonthDayNano)'))", a(2))
}
} else {
unreachable!("handled above")
}
}
// system
"IFNULL" => plain("coalesce"),
"DATABASE" => "'skade'".to_string(),
"USER" => "''".to_string(),
"CONVERT" => {
// {fn CONVERT(x, SQL_VARCHAR)} → CAST(x AS VARCHAR)
let target = a(1).trim().to_ascii_uppercase();
let target = target.strip_prefix("SQL_").unwrap_or(&target).to_string();
let ty = match target.as_str() {
"BIGINT" => "BIGINT",
"INTEGER" => "INTEGER",
"SMALLINT" => "SMALLINT",
"TINYINT" => "TINYINT",
"DOUBLE" | "FLOAT" => "DOUBLE",
"REAL" => "FLOAT",
"DECIMAL" | "NUMERIC" => "DECIMAL",
"BIT" => "BOOLEAN",
"DATE" | "TYPE_DATE" => "DATE",
"TIME" | "TYPE_TIME" => "TIME",
"TIMESTAMP" | "TYPE_TIMESTAMP" => "TIMESTAMP",
"BINARY" | "VARBINARY" | "LONGVARBINARY" => "BYTEA",
_ => "VARCHAR",
};
format!("CAST({} AS {ty})", a(0))
}
// Unknown: keep the call as written; DataFusion may know it.
_ => plain(&name.to_ascii_lowercase()),
}
}
// ── SELECT TOP n ────────────────────────────────────────────────────────────
/// `SELECT [DISTINCT] TOP n …` → `SELECT [DISTINCT] … LIMIT n`, on the
/// OUTERMOST select only, and only when the statement has no LIMIT already.
/// (Power Query's SQL92 generator emits TOP when `SupportsTop = true`; a
/// server that advertises no TOP gets `LIMIT` instead — but drivers differ,
/// so both are accepted.)
fn rewrite_top(sql: &str) -> String {
let b = sql.as_bytes();
let lower = sql.to_ascii_lowercase();
// Find the first keyword `select` outside strings/quotes/parens.
let mut i = 0usize;
let mut depth = 0i32;
let mut select_at = None;
while i < b.len() {
match b[i] {
b'\'' => i = skip_string(b, i),
b'"' => i = skip_quoted(b, i),
b'(' => {
depth += 1;
i += 1;
}
b')' => {
depth -= 1;
i += 1;
}
_ => {
if depth == 0 && lower[i..].starts_with("select") && word_boundary(b, i, i + 6) {
select_at = Some(i);
break;
}
i += 1;
}
}
}
let Some(s) = select_at else { return sql.to_string() };
let mut j = s + 6;
let mut distinct_end = None;
let after = lower[j..].trim_start();
let ws = j + (lower[j..].len() - after.len());
if after.starts_with("distinct") && word_boundary(b, ws, ws + 8) {
distinct_end = Some(ws + 8);
j = ws + 8;
}
let after = lower[j..].trim_start();
let ws = j + (lower[j..].len() - after.len());
if !(after.starts_with("top") && word_boundary(b, ws, ws + 3)) {
return sql.to_string();
}
// parse n
let n_start = ws + 3;
let rest = &sql[n_start..];
let rest_trim = rest.trim_start();
let n_at = n_start + (rest.len() - rest_trim.len());
let n_len = rest_trim.bytes().take_while(|c| c.is_ascii_digit()).count();
if n_len == 0 {
return sql.to_string();
}
let n = &sql[n_at..n_at + n_len];
let head_end = distinct_end.unwrap_or(s + 6);
let body = sql[n_at + n_len..].trim_start();
let mut out = String::with_capacity(sql.len());
out.push_str(&sql[..head_end]);
out.push(' ');
let body = body.trim_end().trim_end_matches(';').trim_end();
out.push_str(body);
if !lower.contains(" limit ") {
out.push_str(" LIMIT ");
out.push_str(n);
}
out
}
fn word_boundary(b: &[u8], start: usize, end: usize) -> bool {
let before_ok = start == 0 || !(b[start - 1].is_ascii_alphanumeric() || b[start - 1] == b'_');
let after_ok = end >= b.len() || !(b[end].is_ascii_alphanumeric() || b[end] == b'_');
before_ok && after_ok
}
#[cfg(test)]
mod tests {
use super::rewrite;
#[test]
fn plain_sql_is_untouched() {
assert!(matches!(rewrite("SELECT * FROM people"), std::borrow::Cow::Borrowed(_)), "plain SQL must not allocate");
for s in [
"SELECT * FROM people",
r#"SELECT "id", "name" FROM "main"."people" WHERE "name" = 'a{b}c' LIMIT 10"#,
"select count(*) from t where s = '{fn not an escape}'",
] {
assert_eq!(rewrite(s), s);
}
}
#[test]
fn typed_literals() {
assert_eq!(rewrite("WHERE d = {d '2024-01-02'}"), "WHERE d = DATE '2024-01-02'");
assert_eq!(rewrite("WHERE t = {t '12:34:56'}"), "WHERE t = TIME '12:34:56'");
assert_eq!(rewrite("WHERE ts >= {ts '2024-01-02 12:34:56.123'}"), "WHERE ts >= TIMESTAMP '2024-01-02 12:34:56.123'");
}
#[test]
fn scalar_functions() {
assert_eq!(rewrite("SELECT {fn UCASE(\"name\")} FROM t"), "SELECT upper(\"name\") FROM t");
assert_eq!(rewrite("SELECT {fn LENGTH({fn LCASE(name)})} FROM t"), "SELECT character_length(lower(name)) FROM t");
assert_eq!(rewrite("WHERE {fn YEAR(ts)} = 2024"), "WHERE date_part('year', ts) = 2024");
assert_eq!(rewrite("SELECT {fn CONVERT(score, SQL_VARCHAR)}"), "SELECT CAST(score AS VARCHAR)");
assert_eq!(rewrite("SELECT {fn IFNULL(score, 0)}"), "SELECT coalesce(score, 0)");
assert_eq!(rewrite("SELECT {fn LOCATE('a', name)}"), "SELECT strpos(name, 'a')");
assert_eq!(rewrite("SELECT {fn TIMESTAMPADD(SQL_TSI_DAY, 3, ts)}"), "SELECT (ts + INTERVAL '3 day')");
assert_eq!(rewrite("SELECT {fn MOD(id, 2)}"), "SELECT ((id) % (2))");
assert_eq!(
rewrite("{fn TIMESTAMPDIFF(SQL_TSI_DAY, {d '2000-01-01'}, ts)}"),
"((to_unixtime(ts) - to_unixtime(DATE '2000-01-01')) / 86400)"
);
// Power Query's Date.From shape: nested escapes with spaces inside the braces.
assert_eq!(
rewrite("{ fn timestampadd(SQL_TSI_DAY, { fn timestampdiff(SQL_TSI_DAY, {d '2000-01-01'}, \"C1\") }, {ts '2000-01-01 00:00:00'}) }"),
"(TIMESTAMP '2000-01-01 00:00:00' + arrow_cast(concat(cast((((to_unixtime(\"C1\") - to_unixtime(DATE '2000-01-01')) / 86400)) as varchar), ' day'), 'Interval(MonthDayNano)'))"
);
assert_eq!(rewrite("SELECT {fn CURDATE()}"), "SELECT current_date()");
assert_eq!(rewrite("SELECT {fn DAYOFWEEK(d)}"), "SELECT (date_part('dow', d) + 1)");
}
#[test]
fn outer_join_and_escape() {
assert_eq!(
rewrite("SELECT * FROM {oj a LEFT OUTER JOIN b ON a.id = b.id}"),
"SELECT * FROM a LEFT OUTER JOIN b ON a.id = b.id"
);
assert_eq!(rewrite("WHERE s LIKE 'a\\%' {escape '\\'}"), "WHERE s LIKE 'a\\%' ESCAPE '\\'");
}
#[test]
fn select_top() {
assert_eq!(rewrite("SELECT TOP 1000 * FROM \"people\""), "SELECT * FROM \"people\" LIMIT 1000");
assert_eq!(rewrite("select top 5 id from t order by id"), "select id from t order by id LIMIT 5");
assert_eq!(rewrite("SELECT DISTINCT TOP 3 name FROM t"), "SELECT DISTINCT name FROM t LIMIT 3");
assert_eq!(rewrite("SELECT TOP 2 * FROM t;"), "SELECT * FROM t LIMIT 2");
// A subquery's TOP is left alone (outermost only).
assert_eq!(rewrite("SELECT * FROM (SELECT TOP 2 * FROM t) q"), "SELECT * FROM (SELECT TOP 2 * FROM t) q");
assert_eq!(rewrite("SELECT toplevel FROM t"), "SELECT toplevel FROM t");
}
#[test]
fn non_ascii_literal_survives() {
assert_eq!(rewrite("WHERE city = 'Umeå' AND d = {d '2024-01-02'}"), "WHERE city = 'Umeå' AND d = DATE '2024-01-02'");
}
}