From 7ce0addb61092058100eb6cf2a9fabd667d0db0e Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sun, 28 Jun 2026 12:15:26 +0800 Subject: [PATCH] fix(sqlite): add if and unistr compatibility functions --- crates/dbx-core/Cargo.toml | 2 +- crates/dbx-core/src/db/sqlite.rs | 197 ++++++++++++++++++++++++++++++- 2 files changed, 197 insertions(+), 2 deletions(-) diff --git a/crates/dbx-core/Cargo.toml b/crates/dbx-core/Cargo.toml index 7d7e914a58..27aff4f428 100644 --- a/crates/dbx-core/Cargo.toml +++ b/crates/dbx-core/Cargo.toml @@ -27,7 +27,7 @@ tokio-postgres = { version = "0.7", features = ["with-chrono-0_4", "with-uuid-1" deadpool-postgres = { version = "0.14", features = ["rt_tokio_1"] } tokio-postgres-rustls = "0.13" webpki-roots = "0.26" -rusqlite = { version = "0.32", features = ["bundled", "load_extension", "backup"] } +rusqlite = { version = "0.32", features = ["bundled", "load_extension", "backup", "functions"] } mysql_async = { version = "0.36", default-features = false, features = ["default-rustls", "client_ed25519", "chrono", "rust_decimal"] } sqlparser = { version = "0.62.0", features = ["visitor"] } redis = { version = "0.32.2", features = ["tokio-comp", "tls-rustls", "tls-rustls-insecure", "tokio-rustls-comp", "sentinel", "cluster-async"] } diff --git a/crates/dbx-core/src/db/sqlite.rs b/crates/dbx-core/src/db/sqlite.rs index 7243f42401..2eda345033 100644 --- a/crates/dbx-core/src/db/sqlite.rs +++ b/crates/dbx-core/src/db/sqlite.rs @@ -1,5 +1,6 @@ use percent_encoding::percent_decode_str; -use rusqlite::types::ValueRef; +use rusqlite::functions::{Context, FunctionFlags}; +use rusqlite::types::{Value, ValueRef}; use rusqlite::{Connection, LoadExtensionGuard, OpenFlags}; use std::collections::HashSet; use std::io::Read; @@ -106,10 +107,179 @@ fn open_sqlite_handle( conn.busy_timeout(std::time::Duration::from_secs(10)).map_err(|e| e.to_string())?; load_sqlite_extensions(&conn, &extensions)?; + register_sqlite_compat_functions(&conn)?; Ok(SqliteHandle { conn: Arc::new(Mutex::new(conn)) }) } +fn register_sqlite_compat_functions(conn: &Connection) -> Result<(), String> { + let flags = FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_DETERMINISTIC | FunctionFlags::SQLITE_INNOCUOUS; + + if !sqlite_function_available(conn, "SELECT if(1, 2, 3)") { + conn.create_scalar_function("if", -1, flags, sqlite_if) + .map_err(|e| format!("SQLite compatibility function registration failed (if): {e}"))?; + } + if !sqlite_function_available(conn, "SELECT unistr('')") { + conn.create_scalar_function("unistr", 1, flags, sqlite_unistr) + .map_err(|e| format!("SQLite compatibility function registration failed (unistr): {e}"))?; + } + + Ok(()) +} + +fn sqlite_function_available(conn: &Connection, sql: &str) -> bool { + conn.query_row(sql, [], |_| Ok(())).is_ok() +} + +fn sqlite_if(ctx: &Context<'_>) -> rusqlite::Result { + if ctx.len() < 2 { + return Err(sqlite_function_error("if() requires at least two arguments")); + } + + let mut i = 0; + while i + 1 < ctx.len() { + if sqlite_truthy(ctx.get_raw(i)) { + return Ok(sqlite_value_ref_to_owned(ctx.get_raw(i + 1))); + } + i += 2; + } + + if ctx.len() % 2 == 1 { + Ok(sqlite_value_ref_to_owned(ctx.get_raw(ctx.len() - 1))) + } else { + Ok(Value::Null) + } +} + +fn sqlite_unistr(ctx: &Context<'_>) -> rusqlite::Result { + let input = match ctx.get_raw(0) { + ValueRef::Null => return Ok(Value::Null), + ValueRef::Integer(value) => value.to_string(), + ValueRef::Real(value) => value.to_string(), + ValueRef::Text(value) | ValueRef::Blob(value) => String::from_utf8_lossy(value).into_owned(), + }; + + sqlite_unistr_text(&input).map(Value::Text) +} + +fn sqlite_truthy(value: ValueRef<'_>) -> bool { + match value { + ValueRef::Null => false, + ValueRef::Integer(value) => value != 0, + ValueRef::Real(value) => value != 0.0, + ValueRef::Text(value) | ValueRef::Blob(value) => sqlite_text_numeric_truthy(&String::from_utf8_lossy(value)), + } +} + +fn sqlite_text_numeric_truthy(text: &str) -> bool { + let trimmed = text.trim_start(); + for end in (1..=trimmed.len()).rev() { + if trimmed.is_char_boundary(end) { + if let Ok(value) = trimmed[..end].parse::() { + return value != 0.0; + } + } + } + false +} + +fn sqlite_value_ref_to_owned(value: ValueRef<'_>) -> Value { + match value { + ValueRef::Null => Value::Null, + ValueRef::Integer(value) => Value::Integer(value), + ValueRef::Real(value) => Value::Real(value), + ValueRef::Text(value) => Value::Text(String::from_utf8_lossy(value).into_owned()), + ValueRef::Blob(value) => Value::Blob(value.to_vec()), + } +} + +fn sqlite_unistr_text(input: &str) -> rusqlite::Result { + let chars: Vec = input.chars().collect(); + let mut result = String::with_capacity(input.len()); + let mut i = 0; + + while i < chars.len() { + if chars[i] != '\\' { + result.push(chars[i]); + i += 1; + continue; + } + + if i + 1 >= chars.len() { + result.push('\\'); + i += 1; + continue; + } + + match chars[i + 1] { + '\\' => { + result.push('\\'); + i += 2; + } + '+' => { + if let Some(ch) = sqlite_unistr_codepoint(&chars, i + 2, 6)? { + result.push(ch); + i += 8; + } else { + result.push('\\'); + i += 1; + } + } + 'u' => { + if let Some(ch) = sqlite_unistr_codepoint(&chars, i + 2, 4)? { + result.push(ch); + i += 6; + } else { + result.push('\\'); + i += 1; + } + } + 'U' => { + if let Some(ch) = sqlite_unistr_codepoint(&chars, i + 2, 8)? { + result.push(ch); + i += 10; + } else { + result.push('\\'); + i += 1; + } + } + _ => { + if let Some(ch) = sqlite_unistr_codepoint(&chars, i + 1, 4)? { + result.push(ch); + i += 5; + } else { + result.push('\\'); + i += 1; + } + } + } + } + + Ok(result) +} + +fn sqlite_unistr_codepoint(chars: &[char], start: usize, digits: usize) -> rusqlite::Result> { + if start + digits > chars.len() { + return Ok(None); + } + + let mut value = 0_u32; + for ch in &chars[start..start + digits] { + let Some(digit) = ch.to_digit(16) else { + return Ok(None); + }; + value = (value << 4) | digit; + } + + std::char::from_u32(value) + .map(Some) + .ok_or_else(|| sqlite_function_error(format!("invalid Unicode codepoint: {value:#X}"))) +} + +fn sqlite_function_error(message: impl Into) -> rusqlite::Error { + rusqlite::Error::UserFunctionError(Box::new(std::io::Error::new(std::io::ErrorKind::InvalidInput, message.into()))) +} + pub fn path_has_sqlite_header(path: &Path) -> Result { let mut file = std::fs::File::open(path).map_err(|e| format!("failed to open file: {e}"))?; let mut header = [0_u8; 16]; @@ -330,6 +500,14 @@ mod tests { assert_eq!(normalize_sqlite_sql("SELECT if (1, 'a', 'b')"), "SELECT IIF (1, 'a', 'b')"); } + #[test] + fn sqlite_unistr_decodes_documented_escapes() { + assert_eq!( + sqlite_unistr_text(r"a\0041\u0042\+000043\U00000044\\z").expect("decode unistr escapes"), + r"aABCD\z" + ); + } + #[tokio::test] async fn view_with_if_function_works_after_normalization() { let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); @@ -349,6 +527,23 @@ mod tests { assert_eq!(result.rows[1][1], serde_json::json!("big")); } + #[tokio::test] + async fn view_with_stored_if_and_unistr_functions_can_be_described_and_queried() { + let pool = connect_path(":memory:").await.expect("connect in-memory SQLite"); + pool.with_connection(|conn| { + conn.execute_batch("CREATE VIEW a AS SELECT if(1, unistr('hello'), 'world') AS a;") + .map_err(|e| e.to_string()) + }) + .expect("create view with original SQLite 3.50 functions"); + + let columns = get_columns(&pool, "", "a").await.expect("describe view"); + assert_eq!(columns.len(), 1); + assert_eq!(columns[0].name, "a"); + + let result = execute_query(&pool, "SELECT * FROM a").await.expect("query view"); + assert_eq!(result.rows[0][0], serde_json::json!("hello")); + } + #[tokio::test] async fn if_rewrite_works_in_direct_query() { let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");