dbx/crates/dbx-core/src/db/sqlite.rs

1491 lines
53 KiB
Rust

use percent_encoding::percent_decode_str;
use rusqlite::types::ValueRef;
use rusqlite::{Connection, LoadExtensionGuard, OpenFlags};
use std::collections::HashSet;
use std::io::Read;
use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Instant;
use super::file_validator::validate_file_path;
use crate::sql::starts_with_executable_sql_keyword;
use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryResult, TableInfo, TriggerInfo};
const SQLITE_DATABASE_HEADER: &[u8; 16] = b"SQLite format 3\0";
#[derive(Clone)]
pub struct SqliteHandle {
conn: Arc<Mutex<Connection>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SqliteExtensionSpec {
pub path: String,
pub entry_point: Option<String>,
}
impl SqliteHandle {
pub fn with_connection<T, F>(&self, f: F) -> Result<T, String>
where
F: FnOnce(&mut Connection) -> Result<T, String>,
{
let mut conn = self.conn.lock().map_err(|e| e.to_string())?;
f(&mut conn)
}
}
pub async fn connect_path(path: &str) -> Result<SqliteHandle, String> {
connect_path_with_options(path, false, Vec::new()).await
}
pub async fn connect_path_with_extensions(
path: &str,
extensions: Vec<SqliteExtensionSpec>,
) -> Result<SqliteHandle, String> {
connect_path_with_options(path, false, extensions).await
}
pub async fn connect_path_create_if_missing(path: &str) -> Result<SqliteHandle, String> {
connect_path_with_options(path, true, Vec::new()).await
}
pub async fn connect_path_create_if_missing_with_extensions(
path: &str,
extensions: Vec<SqliteExtensionSpec>,
) -> Result<SqliteHandle, String> {
connect_path_with_options(path, true, extensions).await
}
async fn connect_path_with_options(
path: &str,
create_if_missing: bool,
extensions: Vec<SqliteExtensionSpec>,
) -> Result<SqliteHandle, String> {
let path = path.to_string();
tokio::task::spawn_blocking(move || open_sqlite_handle(&path, create_if_missing, extensions))
.await
.map_err(|e| e.to_string())?
}
fn open_sqlite_handle(
path: &str,
create_if_missing: bool,
extensions: Vec<SqliteExtensionSpec>,
) -> Result<SqliteHandle, String> {
let is_memory = is_memory_database_path(path);
if !is_memory && !create_if_missing {
validate_file_path(path, is_network_path)?;
}
if !is_memory && create_if_missing {
ensure_parent_dir(path)?;
}
if !is_memory && !is_network_path(path) {
validate_existing_sqlite_file(path)?;
}
let conn = if is_memory {
Connection::open_in_memory().map_err(|e| format!("SQLite connection failed: {e}"))?
} else {
let mut flags = OpenFlags::SQLITE_OPEN_READ_WRITE;
if create_if_missing {
flags |= OpenFlags::SQLITE_OPEN_CREATE;
}
if is_network_path(path) {
flags |= OpenFlags::SQLITE_OPEN_URI;
Connection::open_with_flags(format!("file:{}?vfs=unix-nolock", path), flags)
.map_err(|e| format!("SQLite connection failed: {e}"))?
} else {
Connection::open_with_flags(path, flags).map_err(|e| format!("SQLite connection failed: {e}"))?
}
};
conn.busy_timeout(std::time::Duration::from_secs(10)).map_err(|e| e.to_string())?;
load_sqlite_extensions(&conn, &extensions)?;
Ok(SqliteHandle { conn: Arc::new(Mutex::new(conn)) })
}
pub fn path_has_sqlite_header(path: &Path) -> Result<bool, String> {
let mut file = std::fs::File::open(path).map_err(|e| format!("failed to open file: {e}"))?;
let mut header = [0_u8; 16];
match file.read_exact(&mut header) {
Ok(()) => Ok(&header == SQLITE_DATABASE_HEADER),
Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => Ok(false),
Err(e) => Err(format!("failed to read file header: {e}")),
}
}
fn validate_existing_sqlite_file(path: &str) -> Result<(), String> {
let path = Path::new(path);
if !path.exists() {
return Ok(());
}
let metadata = path.metadata().map_err(|e| format!("failed to inspect SQLite database file: {e}"))?;
if metadata.len() == 0 {
return Ok(());
}
if path_has_sqlite_header(path)? {
return Ok(());
}
Err("Selected file is not a valid SQLite database file.".to_string())
}
fn load_sqlite_extensions(conn: &Connection, extensions: &[SqliteExtensionSpec]) -> Result<(), String> {
if extensions.is_empty() {
return Ok(());
}
// Extension loading is enabled only for the trusted paths from the connection config.
let _guard =
unsafe { LoadExtensionGuard::new(conn) }.map_err(|e| format!("SQLite extension loading failed: {e}"))?;
for extension in extensions {
unsafe { conn.load_extension(&extension.path, extension.entry_point.as_deref()) }
.map_err(|e| format!("SQLite extension load failed ({}): {e}", extension.path))?;
}
Ok(())
}
pub fn sqlite_extension_specs_from_url_params(params: Option<&str>) -> Vec<SqliteExtensionSpec> {
params
.unwrap_or("")
.trim()
.trim_start_matches('?')
.split('&')
.filter_map(|part| {
let (raw_key, raw_value) = part.split_once('=').unwrap_or((part, ""));
let key = decode_url_param(raw_key);
if key != "sqlite_extension" && key != "sqlite_extensions" {
return None;
}
Some(decode_url_param(raw_value))
})
.flat_map(|value| value.lines().filter_map(parse_sqlite_extension_spec).collect::<Vec<_>>())
.collect()
}
fn parse_sqlite_extension_spec(value: &str) -> Option<SqliteExtensionSpec> {
let value = value.trim();
if value.is_empty() {
return None;
}
let (path, entry_point) = match value.rsplit_once('|') {
Some((path, entry_point)) if !path.trim().is_empty() && !entry_point.trim().is_empty() => {
(path.trim(), Some(entry_point.trim().to_string()))
}
_ => (value, None),
};
Some(SqliteExtensionSpec { path: path.to_string(), entry_point })
}
fn decode_url_param(value: &str) -> String {
percent_decode_str(&value.replace('+', " ")).decode_utf8_lossy().into_owned()
}
fn ensure_parent_dir(path: &str) -> Result<(), String> {
if let Some(parent) = Path::new(path).parent() {
if !parent.as_os_str().is_empty() {
std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
}
}
Ok(())
}
fn is_network_path(path: &str) -> bool {
path.starts_with("\\\\") || path.starts_with("//") || path.contains("wsl.localhost") || path.contains("wsl$")
}
pub fn is_memory_database_path(path: &str) -> bool {
path.trim().eq_ignore_ascii_case(":memory:")
}
#[cfg(test)]
#[allow(clippy::items_after_test_module)]
mod tests {
use super::*;
#[tokio::test]
async fn connect_path_supports_memory_database_across_statements() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(&pool, "CREATE TABLE memory_probe (id INTEGER PRIMARY KEY, name TEXT);")
.await
.expect("create table");
execute_query(&pool, "INSERT INTO memory_probe (name) VALUES ('Ada');").await.expect("insert row");
let result = execute_query(&pool, "SELECT name FROM memory_probe WHERE id = 1;").await.expect("select row");
assert_eq!(result.rows[0][0], serde_json::json!("Ada"));
}
#[tokio::test]
async fn create_if_missing_rejects_existing_non_sqlite_file() {
let path = std::env::temp_dir().join(format!("dbx-not-sqlite-{}.png", uuid::Uuid::new_v4()));
std::fs::write(&path, b"\x89PNG\r\n\x1a\nnot sqlite").unwrap();
let err = match connect_path_create_if_missing(path.to_str().unwrap()).await {
Ok(_) => panic!("non-SQLite file should be rejected"),
Err(err) => err,
};
assert!(err.contains("not a valid SQLite database"));
let _ = std::fs::remove_file(path);
}
#[tokio::test]
async fn create_if_missing_allows_empty_custom_suffix_file() {
let path = std::env::temp_dir().join(format!("dbx-empty-sqlite-{}.conf", uuid::Uuid::new_v4()));
std::fs::write(&path, b"").unwrap();
let pool = connect_path_create_if_missing(path.to_str().unwrap()).await.expect("empty file can become SQLite");
execute_query(&pool, "CREATE TABLE t (id INTEGER);").await.expect("write sqlite schema");
let _ = std::fs::remove_file(path);
}
#[tokio::test]
async fn create_if_missing_allows_sqlite_database_with_custom_suffix() {
let path = std::env::temp_dir().join(format!("dbx-custom-sqlite-{}.conf", uuid::Uuid::new_v4()));
{
let pool = connect_path_create_if_missing(path.to_str().unwrap()).await.expect("create sqlite");
execute_query(&pool, "CREATE TABLE t (id INTEGER);").await.expect("write sqlite schema");
}
let reopened = connect_path_create_if_missing(path.to_str().unwrap()).await.expect("reopen sqlite");
let result = execute_query(&reopened, "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 't';")
.await
.expect("query sqlite schema");
assert_eq!(result.rows[0][0], serde_json::json!("t"));
let _ = std::fs::remove_file(path);
}
#[test]
fn sqlite_extension_specs_parse_repeated_and_multiline_url_params() {
let params = "cache=shared&sqlite_extension=%2Fopt%2Fregexp.dylib&sqlite_extensions=%2Fopt%2Ftext.dylib%7Csqlite3_text_init%0A%2Fopt%2Fcrypto.dylib";
assert_eq!(
sqlite_extension_specs_from_url_params(Some(params)),
vec![
SqliteExtensionSpec { path: "/opt/regexp.dylib".to_string(), entry_point: None },
SqliteExtensionSpec {
path: "/opt/text.dylib".to_string(),
entry_point: Some("sqlite3_text_init".to_string()),
},
SqliteExtensionSpec { path: "/opt/crypto.dylib".to_string(), entry_point: None },
],
);
}
#[test]
fn sqlite_extension_specs_ignore_empty_values() {
assert!(sqlite_extension_specs_from_url_params(Some("sqlite_extension=&sqlite_extensions=%0A")).is_empty());
}
#[test]
fn normalize_if_to_iif_basic() {
assert_eq!(normalize_sqlite_sql("SELECT if(1, 'a', 'b')"), "SELECT IIF(1, 'a', 'b')");
assert_eq!(normalize_sqlite_sql("SELECT if(1, if(0, 'x', 'y'), 'b')"), "SELECT IIF(1, IIF(0, 'x', 'y'), 'b')");
}
#[test]
fn normalize_substring_to_substr() {
assert_eq!(normalize_sqlite_sql("SELECT substring(name, 1, 3) FROM t"), "SELECT substr(name, 1, 3) FROM t");
assert_eq!(normalize_sqlite_sql("SELECT substring(name, 2) FROM t"), "SELECT substr(name, 2) FROM t");
}
#[test]
fn normalize_preserves_string_literals() {
let sql = "SELECT 'if(1,2,3)' AS literal, 'substring(x,1,2)', if(1, 'ok', 'no')";
let normalized = normalize_sqlite_sql(sql);
assert_eq!(normalized, "SELECT 'if(1,2,3)' AS literal, 'substring(x,1,2)', IIF(1, 'ok', 'no')");
}
#[test]
fn normalize_preserves_line_comments() {
let sql = "-- if(1,2,3) is a comment\nSELECT if(1, 'x', 'y')";
let normalized = normalize_sqlite_sql(sql);
assert_eq!(normalized, "-- if(1,2,3) is a comment\nSELECT IIF(1, 'x', 'y')");
}
#[test]
fn normalize_preserves_block_comments() {
let sql = "/* if(1,2,3) */ SELECT if(1, 'x', 'y')";
let normalized = normalize_sqlite_sql(sql);
assert_eq!(normalized, "/* if(1,2,3) */ SELECT IIF(1, 'x', 'y')");
}
#[test]
fn normalize_does_not_match_inside_words() {
let sql = "SELECT difference, stiff, ifsubstring FROM t";
let normalized = normalize_sqlite_sql(sql);
assert_eq!(normalized, sql);
}
#[test]
fn normalize_if_with_spaces_before_paren() {
assert_eq!(normalize_sqlite_sql("SELECT if (1, 'a', 'b')"), "SELECT IIF (1, 'a', 'b')");
}
#[tokio::test]
async fn view_with_if_function_works_after_normalization() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(&pool, "CREATE TABLE t (x INTEGER); INSERT INTO t VALUES (1), (2), (3);")
.await
.expect("create and populate table");
execute_query(&pool, "CREATE VIEW v AS SELECT x, IIF(x > 1, 'big', 'small') AS label FROM t")
.await
.expect("create view");
let result = execute_query(&pool, "SELECT * FROM v ORDER BY x").await.expect("query view");
assert_eq!(result.rows.len(), 3);
assert_eq!(result.rows[0][1], serde_json::json!("small"));
assert_eq!(result.rows[1][1], serde_json::json!("big"));
}
#[tokio::test]
async fn if_rewrite_works_in_direct_query() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
let result = execute_query(&pool, "SELECT if(1 = 1, 'yes', 'no') AS answer")
.await
.expect("if() should be rewritten to IIF()");
assert_eq!(result.rows[0][0], serde_json::json!("yes"));
}
#[tokio::test]
async fn bundled_sqlite_math_functions_are_available() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
let floor_result =
execute_query(&pool, "WITH test(x) AS (VALUES (1.1), (1.2), (1.3)) SELECT FLOOR(x) FROM test")
.await
.expect("FLOOR() should be available");
assert_eq!(floor_result.rows.len(), 3);
for row in floor_result.rows {
assert_eq!(row[0].as_f64(), Some(1.0));
}
let result = execute_query(&pool, "SELECT ACOS(1.0), ACOSH(1.0), ASIN(0.0), CEIL(1.2), PI()")
.await
.expect("SQLite math functions should be available");
assert_eq!(result.rows[0][0].as_f64(), Some(0.0));
assert_eq!(result.rows[0][1].as_f64(), Some(0.0));
assert_eq!(result.rows[0][2].as_f64(), Some(0.0));
assert_eq!(result.rows[0][3].as_f64(), Some(2.0));
let pi = result.rows[0][4].as_f64().expect("PI() returns a real value");
assert!((std::f64::consts::PI - pi).abs() < 0.00001);
}
#[tokio::test]
async fn substring_rewrite_works_in_direct_query() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(&pool, "CREATE TABLE t (name TEXT); INSERT INTO t VALUES ('hello');").await.expect("setup");
let result = execute_query(&pool, "SELECT substring(name, 1, 2) AS s FROM t")
.await
.expect("substring() should be rewritten to substr()");
assert_eq!(result.rows[0][0], serde_json::json!("he"));
}
#[tokio::test]
async fn both_rewrites_combined() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(&pool, "CREATE TABLE t (x INTEGER); INSERT INTO t VALUES (1), (2);").await.expect("setup");
let result = execute_query(&pool, "SELECT substring(if(x > 1, 'big', 'small'), 1, 1) AS s FROM t ORDER BY x")
.await
.expect("combined rewrite");
assert_eq!(result.rows[0][0], serde_json::json!("s"));
assert_eq!(result.rows[1][0], serde_json::json!("b"));
}
fn parse_pk(sql: &str) -> Vec<String> {
let mut cols: Vec<String> = parse_sqlite_autoincrement_pk_columns(sql).into_iter().collect();
cols.sort();
cols
}
#[test]
fn parses_implicit_integer_primary_key_as_autoincrement() {
assert_eq!(parse_pk("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT)"), vec!["id".to_string()]);
}
#[test]
fn parses_explicit_integer_primary_key_autoincrement() {
assert_eq!(
parse_pk("CREATE TABLE t (id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT)"),
vec!["id".to_string()]
);
}
#[test]
fn parses_ef_core_style_named_constraint_primary_key_autoincrement() {
// The actual table from issue #1129.
let sql = r#"CREATE TABLE "OnlineLogs" (
"OnlineLogId" INTEGER NOT NULL CONSTRAINT "PK_OnlineLogs" PRIMARY KEY AUTOINCREMENT,
"LogTime" TEXT NOT NULL,
"ReportedAddresses" TEXT NOT NULL,
"DeviceId" TEXT NOT NULL
)"#;
assert_eq!(parse_pk(sql), vec!["onlinelogid".to_string()]);
}
#[test]
fn does_not_match_non_integer_primary_key() {
assert!(parse_sqlite_autoincrement_pk_columns("CREATE TABLE t (id INT PRIMARY KEY, name TEXT)").is_empty());
assert!(parse_sqlite_autoincrement_pk_columns("CREATE TABLE t (id BIGINT PRIMARY KEY, name TEXT)").is_empty());
}
#[test]
fn does_not_match_without_rowid_table() {
let sql = "CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT) WITHOUT ROWID";
assert!(parse_sqlite_autoincrement_pk_columns(sql).is_empty());
}
#[test]
fn does_not_match_composite_primary_key() {
let sql = "CREATE TABLE t (a INTEGER, b INTEGER, PRIMARY KEY (a, b))";
assert!(parse_sqlite_autoincrement_pk_columns(sql).is_empty());
}
#[test]
fn parses_table_level_single_column_primary_key_for_integer() {
let sql = "CREATE TABLE t (id INTEGER NOT NULL, name TEXT, PRIMARY KEY (id))";
assert_eq!(parse_pk(sql), vec!["id".to_string()]);
}
#[test]
fn ignores_non_pk_integer_not_null_column() {
let sql = "CREATE TABLE t (id INTEGER PRIMARY KEY, count INTEGER NOT NULL)";
assert_eq!(parse_pk(sql), vec!["id".to_string()]);
}
#[test]
fn parser_falls_back_to_empty_on_garbage_sql() {
assert!(parse_sqlite_autoincrement_pk_columns("not a create table statement").is_empty());
assert!(parse_sqlite_autoincrement_pk_columns("").is_empty());
}
#[test]
fn parser_skips_check_expression_with_primary_key_token() {
// PRIMARY KEY tokens inside a CHECK expression must not falsely mark the column.
let sql = r#"CREATE TABLE t (
id INTEGER,
kind TEXT CHECK (kind IN ('PRIMARY KEY', 'OTHER')),
PRIMARY KEY (id)
)"#;
assert_eq!(parse_pk(sql), vec!["id".to_string()]);
}
#[test]
fn parser_handles_block_and_line_comments() {
let sql = r#"CREATE TABLE t (
-- line comment with INTEGER PRIMARY KEY tokens
/* block comment INTEGER PRIMARY KEY */
id INTEGER PRIMARY KEY,
name TEXT
)"#;
assert_eq!(parse_pk(sql), vec!["id".to_string()]);
}
#[tokio::test]
async fn get_columns_marks_integer_primary_key_as_autoincrement() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(&pool, "CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT NOT NULL)").await.expect("create");
let cols = get_columns(&pool, "main", "t").await.expect("get_columns");
let id = cols.iter().find(|c| c.name == "id").expect("id col");
assert_eq!(id.extra.as_deref(), Some("autoincrement"));
let name = cols.iter().find(|c| c.name == "name").expect("name col");
assert!(name.extra.is_none());
}
#[tokio::test]
async fn get_columns_marks_ef_core_style_autoincrement_primary_key() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(
&pool,
r#"CREATE TABLE "OnlineLogs" (
"OnlineLogId" INTEGER NOT NULL CONSTRAINT "PK_OnlineLogs" PRIMARY KEY AUTOINCREMENT,
"LogTime" TEXT NOT NULL,
"DeviceId" TEXT NOT NULL
)"#,
)
.await
.expect("create");
let cols = get_columns(&pool, "main", "OnlineLogs").await.expect("get_columns");
let id = cols.iter().find(|c| c.name == "OnlineLogId").expect("OnlineLogId");
assert_eq!(id.extra.as_deref(), Some("autoincrement"));
for other in cols.iter().filter(|c| c.name != "OnlineLogId") {
assert!(other.extra.is_none(), "{} should not be autoincrement", other.name);
}
}
#[tokio::test]
async fn get_columns_skips_without_rowid_table() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(&pool, "CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT NOT NULL) WITHOUT ROWID")
.await
.expect("create");
let cols = get_columns(&pool, "main", "t").await.expect("get_columns");
let id = cols.iter().find(|c| c.name == "id").expect("id col");
assert!(id.extra.is_none());
}
#[tokio::test]
async fn get_columns_skips_composite_primary_key() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
execute_query(&pool, "CREATE TABLE t (a INTEGER NOT NULL, b INTEGER NOT NULL, PRIMARY KEY (a, b))")
.await
.expect("create");
let cols = get_columns(&pool, "main", "t").await.expect("get_columns");
for col in &cols {
assert!(col.extra.is_none(), "{} should not be autoincrement", col.name);
}
}
#[tokio::test]
async fn get_columns_skips_non_integer_primary_key() {
let pool = connect_path(":memory:").await.expect("connect in-memory SQLite");
// Use BIGINT to avoid SQLite's strict-table parser quirks; INT is sometimes promoted in SQLite.
execute_query(&pool, "CREATE TABLE t (id BIGINT PRIMARY KEY, name TEXT)").await.expect("create");
let cols = get_columns(&pool, "main", "t").await.expect("get_columns");
let id = cols.iter().find(|c| c.name == "id").expect("id col");
assert!(id.extra.is_none());
}
}
pub async fn list_databases(_pool: &SqliteHandle) -> Result<Vec<DatabaseInfo>, String> {
Ok(vec![DatabaseInfo { name: "main".to_string() }])
}
pub async fn list_tables(pool: &SqliteHandle, _schema: &str) -> Result<Vec<TableInfo>, String> {
let pool = pool.clone();
tokio::task::spawn_blocking(move || {
pool.with_connection(|conn| {
let mut stmt = conn
.prepare(
"SELECT name, type FROM sqlite_master \
WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name",
)
.map_err(|e| e.to_string())?;
let rows = stmt
.query_map([], |row| {
let table_type: String = row.get(1)?;
Ok(TableInfo {
name: row.get(0)?,
table_type: if table_type == "view" { "VIEW".to_string() } else { "BASE TABLE".to_string() },
comment: None,
parent_schema: None,
parent_name: None,
})
})
.map_err(|e| e.to_string())?;
rows.collect::<Result<Vec<_>, _>>().map_err(|e| e.to_string())
})
})
.await
.map_err(|e| e.to_string())?
}
pub async fn get_columns(pool: &SqliteHandle, _schema: &str, table: &str) -> Result<Vec<ColumnInfo>, String> {
let pool = pool.clone();
let table = table.to_string();
tokio::task::spawn_blocking(move || {
let sql = format!("PRAGMA table_info(\"{}\")", table.replace('"', "\"\""));
pool.with_connection(|conn| {
let autoincrement_columns = sqlite_autoincrement_pk_columns(conn, &table).unwrap_or_default();
let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?;
let rows = stmt
.query_map([], |row| {
let name: String = row.get("name")?;
let is_pk = row.get::<_, i32>("pk")? > 0;
let extra = if is_pk && autoincrement_columns.contains(&name.to_ascii_lowercase()) {
Some("autoincrement".to_string())
} else {
None
};
Ok(ColumnInfo {
name,
data_type: row.get("type")?,
is_nullable: row.get::<_, i32>("notnull")? == 0,
column_default: row.get("dflt_value")?,
is_primary_key: is_pk,
extra,
comment: None,
numeric_precision: None,
numeric_scale: None,
character_maximum_length: None,
})
})
.map_err(|e| e.to_string())?;
rows.collect::<Result<Vec<_>, _>>().map_err(|e| e.to_string())
})
})
.await
.map_err(|e| e.to_string())?
}
/// Read `sqlite_master.sql` for `table` and return the lowercase column names that
/// are rowid-alias autoincrement primary keys (i.e. SQLite will assign a value when
/// the column is omitted from an INSERT). Returns `None` only on connection / query
/// errors; an unparseable build statement yields `Some(empty)`.
fn sqlite_autoincrement_pk_columns(conn: &Connection, table: &str) -> Option<HashSet<String>> {
let create_sql: Option<String> = conn
.query_row("SELECT sql FROM sqlite_master WHERE type = 'table' AND name = ?1", [table], |row| row.get(0))
.ok()
.flatten();
Some(parse_sqlite_autoincrement_pk_columns(create_sql.as_deref()?))
}
/// Parse a SQLite `CREATE TABLE` statement and return the lowercase names of columns
/// that are rowid-alias autoincrement primary keys.
///
/// A column is recognized when ALL of the following hold:
/// - The table is NOT declared `WITHOUT ROWID`.
/// - The column's declared type, after case-insensitive normalization, is exactly
/// `INTEGER` (NOT `INT`, `BIGINT`, `SMALLINT`, etc.).
/// - The column is the (only) primary key, declared either inline (`PRIMARY KEY`,
/// optionally with `AUTOINCREMENT`) or via a single-column table-level
/// `PRIMARY KEY (col)` constraint.
///
/// On any parse failure (malformed SQL, unrecognized syntax) the function returns
/// an empty set rather than panicking — callers fall back to the conservative
/// behavior of treating the column as a normal NOT NULL column.
fn parse_sqlite_autoincrement_pk_columns(create_sql: &str) -> HashSet<String> {
let body = match extract_create_table_body(create_sql) {
Some(body) => body,
None => return HashSet::new(),
};
if has_without_rowid_clause(&body.tail) {
return HashSet::new();
}
let entries = split_table_body_entries(&body.body);
// First pass: find table-level PRIMARY KEY (col) — a single-column primary key
// that may apply to a column declared as INTEGER elsewhere in the body.
let mut table_level_pk: Option<String> = None;
let mut has_composite_table_pk = false;
for entry in &entries {
if let Some(cols) = parse_table_level_primary_key(entry) {
if cols.len() == 1 {
if table_level_pk.is_none() && !has_composite_table_pk {
table_level_pk = Some(cols.into_iter().next().unwrap());
}
} else if cols.len() > 1 {
has_composite_table_pk = true;
table_level_pk = None;
}
}
}
if has_composite_table_pk {
return HashSet::new();
}
let mut found: HashSet<String> = HashSet::new();
let mut inline_pk_count = 0_usize;
let mut inline_pk_candidate: Option<String> = None;
for entry in &entries {
if parse_table_level_primary_key(entry).is_some() {
continue;
}
if is_table_level_constraint(entry) {
continue;
}
let Some(column) = parse_column_definition(entry) else {
continue;
};
if column.has_inline_pk {
inline_pk_count += 1;
inline_pk_candidate = Some(column.name.clone());
if column.is_integer_type {
found.insert(column.name.clone());
}
}
if let Some(ref pk_name) = table_level_pk {
if pk_name.eq_ignore_ascii_case(&column.name) && column.is_integer_type {
found.insert(column.name.clone());
}
}
}
// Multiple inline PRIMARY KEY columns means a composite key — clear the inline matches.
if inline_pk_count > 1 {
if let Some(name) = inline_pk_candidate {
found.remove(&name);
}
// Also drop any other inline PK columns we may have inserted.
// (Conservative: walk entries again and remove inline PK names that ended up in `found`.)
let mut to_remove: Vec<String> = Vec::new();
for entry in &entries {
if let Some(column) = parse_column_definition(entry) {
if column.has_inline_pk && found.contains(&column.name) {
to_remove.push(column.name);
}
}
}
for name in to_remove {
found.remove(&name);
}
}
found
}
struct CreateTableBody {
body: String,
tail: String,
}
fn extract_create_table_body(create_sql: &str) -> Option<CreateTableBody> {
let stripped = strip_sql_comments(create_sql);
let lower = stripped.to_ascii_lowercase();
if !lower.contains("create") || !lower.contains("table") {
return None;
}
// Find the first top-level '(' after the table name.
let bytes = stripped.as_bytes();
let mut start = None;
for (i, &b) in bytes.iter().enumerate() {
if b == b'(' {
start = Some(i);
break;
}
}
let start = start?;
let mut depth = 0_usize;
let mut end = None;
let mut chars = stripped[start..].char_indices();
while let Some((rel, ch)) = chars.next() {
match ch {
'(' => depth += 1,
')' => {
depth -= 1;
if depth == 0 {
end = Some(start + rel);
break;
}
}
'\'' | '"' | '`' => {
// skip a quoted string / identifier in the loop directly
let quote = ch;
while let Some((_, qch)) = chars.next() {
if qch == quote {
// SQLite supports doubled quote as escape inside identifiers.
// Peek next char without consuming.
let mut peek = chars.clone();
if let Some((_, next_ch)) = peek.next() {
if next_ch == quote {
chars.next();
continue;
}
}
break;
}
}
}
'[' => {
// SQL Server style identifier — closes at first ']'.
for (_, qch) in chars.by_ref() {
if qch == ']' {
break;
}
}
}
_ => {}
}
}
let end = end?;
let body = stripped[start + 1..end].to_string();
let tail = stripped[end + 1..].to_string();
Some(CreateTableBody { body, tail })
}
fn has_without_rowid_clause(tail: &str) -> bool {
let normalized: String = tail.split_whitespace().collect::<Vec<_>>().join(" ").to_ascii_lowercase();
normalized.contains("without rowid")
}
fn strip_sql_comments(input: &str) -> String {
let mut out = String::with_capacity(input.len());
let bytes = input.as_bytes();
let mut i = 0;
while i < bytes.len() {
let b = bytes[i];
if b == b'-' && i + 1 < bytes.len() && bytes[i + 1] == b'-' {
// line comment
i += 2;
while i < bytes.len() && bytes[i] != b'\n' {
i += 1;
}
} else if b == b'/' && i + 1 < bytes.len() && bytes[i + 1] == b'*' {
i += 2;
while i + 1 < bytes.len() && !(bytes[i] == b'*' && bytes[i + 1] == b'/') {
i += 1;
}
if i + 1 < bytes.len() {
i += 2;
} else {
i = bytes.len();
}
} else if b == b'\'' || b == b'"' || b == b'`' {
// copy quoted segment as-is (we still need it for identifier parsing later)
let quote = b;
out.push(b as char);
i += 1;
while i < bytes.len() {
let qb = bytes[i];
out.push(qb as char);
if qb == quote {
if i + 1 < bytes.len() && bytes[i + 1] == quote {
// doubled quote escape
out.push(quote as char);
i += 2;
continue;
}
i += 1;
break;
}
i += 1;
}
} else if b == b'[' {
out.push('[');
i += 1;
while i < bytes.len() {
let qb = bytes[i];
out.push(qb as char);
i += 1;
if qb == b']' {
break;
}
}
} else {
// copy as char (handle multi-byte utf-8 by walking)
out.push(input[i..].chars().next().unwrap());
i += input[i..].chars().next().unwrap().len_utf8();
}
}
out
}
fn split_table_body_entries(body: &str) -> Vec<String> {
let mut entries = Vec::new();
let mut current = String::new();
let mut depth = 0_usize;
let mut chars = body.chars().peekable();
while let Some(ch) = chars.next() {
match ch {
'(' => {
depth += 1;
current.push(ch);
}
')' => {
depth = depth.saturating_sub(1);
current.push(ch);
}
',' if depth == 0 => {
let trimmed = current.trim();
if !trimmed.is_empty() {
entries.push(trimmed.to_string());
}
current.clear();
}
'\'' | '"' | '`' => {
let quote = ch;
current.push(ch);
while let Some(qch) = chars.next() {
current.push(qch);
if qch == quote {
if let Some(&next_ch) = chars.peek() {
if next_ch == quote {
current.push(chars.next().unwrap());
continue;
}
}
break;
}
}
}
'[' => {
current.push(ch);
for qch in chars.by_ref() {
current.push(qch);
if qch == ']' {
break;
}
}
}
_ => current.push(ch),
}
}
let trimmed = current.trim();
if !trimmed.is_empty() {
entries.push(trimmed.to_string());
}
entries
}
struct ColumnDefinition {
name: String,
is_integer_type: bool,
has_inline_pk: bool,
}
fn parse_column_definition(entry: &str) -> Option<ColumnDefinition> {
let mut tokens = tokenize_entry(entry);
if tokens.is_empty() {
return None;
}
// Skip leading "CONSTRAINT name" if it appears (rare in column defs but tolerated).
if tokens[0].kind == TokenKind::Keyword && tokens[0].value.eq_ignore_ascii_case("constraint") && tokens.len() >= 2 {
// not a column definition
return None;
}
let name_token = tokens.remove(0);
if name_token.kind != TokenKind::Identifier {
return None;
}
let name_lower = name_token.value.to_ascii_lowercase();
// Type token: optional, followed by optional parenthesized size.
let mut is_integer_type = false;
if let Some(first) = tokens.first() {
if first.kind == TokenKind::Identifier {
if first.value.eq_ignore_ascii_case("integer") {
is_integer_type = true;
}
// consume the type token; also consume size like "(10, 2)"
tokens.remove(0);
if let Some(t) = tokens.first() {
if t.value == "(" {
// consume balanced parens
let mut depth = 0_usize;
while !tokens.is_empty() {
let t = tokens.remove(0);
if t.value == "(" {
depth += 1;
} else if t.value == ")" {
depth = depth.saturating_sub(1);
if depth == 0 {
break;
}
}
}
}
}
}
}
let has_inline_pk = tokens_contain_primary_key(&tokens);
Some(ColumnDefinition { name: name_lower, is_integer_type, has_inline_pk })
}
fn tokens_contain_primary_key(tokens: &[Token]) -> bool {
for window in tokens.windows(2) {
if window[0].kind == TokenKind::Keyword
&& window[0].value.eq_ignore_ascii_case("primary")
&& window[1].kind == TokenKind::Keyword
&& window[1].value.eq_ignore_ascii_case("key")
{
return true;
}
}
false
}
fn is_table_level_constraint(entry: &str) -> bool {
let trimmed = entry.trim_start();
let lower = trimmed.to_ascii_lowercase();
lower.starts_with("constraint")
|| lower.starts_with("primary key")
|| lower.starts_with("unique")
|| lower.starts_with("check")
|| lower.starts_with("foreign key")
}
/// If the entry is a table-level `PRIMARY KEY (col[, col, ...])` constraint,
/// return the lowercase column names. Otherwise `None`.
fn parse_table_level_primary_key(entry: &str) -> Option<Vec<String>> {
let tokens = tokenize_entry(entry);
let mut idx = 0;
if idx < tokens.len()
&& tokens[idx].kind == TokenKind::Keyword
&& tokens[idx].value.eq_ignore_ascii_case("constraint")
{
idx += 1;
if idx < tokens.len() && tokens[idx].kind == TokenKind::Identifier {
idx += 1;
}
}
if idx + 1 >= tokens.len() {
return None;
}
if !(tokens[idx].kind == TokenKind::Keyword
&& tokens[idx].value.eq_ignore_ascii_case("primary")
&& tokens[idx + 1].kind == TokenKind::Keyword
&& tokens[idx + 1].value.eq_ignore_ascii_case("key"))
{
return None;
}
idx += 2;
if idx >= tokens.len() || tokens[idx].value != "(" {
return None;
}
idx += 1;
let mut cols = Vec::new();
while idx < tokens.len() && tokens[idx].value != ")" {
if tokens[idx].kind == TokenKind::Identifier {
cols.push(tokens[idx].value.to_ascii_lowercase());
}
idx += 1;
// skip optional ASC/DESC and a comma
while idx < tokens.len() && tokens[idx].value != "," && tokens[idx].value != ")" {
idx += 1;
}
if idx < tokens.len() && tokens[idx].value == "," {
idx += 1;
}
}
Some(cols)
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum TokenKind {
Identifier,
Keyword,
Punct,
Other,
}
#[derive(Debug, Clone)]
struct Token {
value: String,
kind: TokenKind,
}
fn tokenize_entry(entry: &str) -> Vec<Token> {
let mut tokens = Vec::new();
let mut chars = entry.chars().peekable();
while let Some(&ch) = chars.peek() {
if ch.is_whitespace() {
chars.next();
continue;
}
if ch == '"' || ch == '`' {
let quote = ch;
chars.next();
let mut value = String::new();
while let Some(&qch) = chars.peek() {
chars.next();
if qch == quote {
if chars.peek() == Some(&quote) {
value.push(quote);
chars.next();
continue;
}
break;
}
value.push(qch);
}
tokens.push(Token { value, kind: TokenKind::Identifier });
continue;
}
if ch == '[' {
chars.next();
let mut value = String::new();
while let Some(&qch) = chars.peek() {
chars.next();
if qch == ']' {
break;
}
value.push(qch);
}
tokens.push(Token { value, kind: TokenKind::Identifier });
continue;
}
if ch == '\'' {
// string literal — skip
chars.next();
while let Some(&qch) = chars.peek() {
chars.next();
if qch == '\'' {
if chars.peek() == Some(&'\'') {
chars.next();
continue;
}
break;
}
}
continue;
}
if ch == '(' || ch == ')' || ch == ',' {
chars.next();
tokens.push(Token { value: ch.to_string(), kind: TokenKind::Punct });
continue;
}
if ch.is_ascii_alphabetic() || ch == '_' {
let mut value = String::new();
while let Some(&wch) = chars.peek() {
if wch.is_ascii_alphanumeric() || wch == '_' {
value.push(wch);
chars.next();
} else {
break;
}
}
let kind = if is_sql_keyword(&value) { TokenKind::Keyword } else { TokenKind::Identifier };
tokens.push(Token { value, kind });
continue;
}
// anything else — skip but record for completeness
chars.next();
tokens.push(Token { value: ch.to_string(), kind: TokenKind::Other });
}
tokens
}
fn is_sql_keyword(value: &str) -> bool {
matches!(
value.to_ascii_lowercase().as_str(),
"constraint"
| "primary"
| "key"
| "not"
| "null"
| "default"
| "unique"
| "check"
| "foreign"
| "references"
| "on"
| "delete"
| "update"
| "cascade"
| "set"
| "restrict"
| "no"
| "action"
| "deferrable"
| "initially"
| "deferred"
| "immediate"
| "match"
| "collate"
| "autoincrement"
| "asc"
| "desc"
| "generated"
| "always"
| "stored"
| "virtual"
| "as"
)
}
pub async fn list_indexes(pool: &SqliteHandle, _schema: &str, table: &str) -> Result<Vec<IndexInfo>, String> {
let pool = pool.clone();
let table = table.to_string();
tokio::task::spawn_blocking(move || {
let safe_table = table.replace('"', "\"\"");
pool.with_connection(|conn| {
let mut stmt = conn.prepare(&format!("PRAGMA index_list(\"{safe_table}\")")).map_err(|e| e.to_string())?;
let idx_rows = stmt
.query_map([], |row| {
Ok((
row.get::<_, String>("name")?,
row.get::<_, i32>("unique")? != 0,
row.get::<_, String>("origin")?,
))
})
.map_err(|e| e.to_string())?
.collect::<Result<Vec<_>, _>>()
.map_err(|e| e.to_string())?;
let mut indexes = Vec::new();
for (name, is_unique, origin) in idx_rows {
let safe_name = name.replace('"', "\"\"");
let mut col_stmt =
conn.prepare(&format!("PRAGMA index_info(\"{safe_name}\")")).map_err(|e| e.to_string())?;
let columns = col_stmt
.query_map([], |row| row.get::<_, String>("name"))
.map_err(|e| e.to_string())?
.collect::<Result<Vec<_>, _>>()
.map_err(|e| e.to_string())?;
indexes.push(IndexInfo {
name,
columns,
is_unique,
is_primary: origin == "pk",
filter: None,
index_type: None,
included_columns: None,
comment: None,
});
}
Ok(indexes)
})
})
.await
.map_err(|e| e.to_string())?
}
pub async fn list_foreign_keys(pool: &SqliteHandle, _schema: &str, table: &str) -> Result<Vec<ForeignKeyInfo>, String> {
let pool = pool.clone();
let table = table.to_string();
tokio::task::spawn_blocking(move || {
let sql = format!("PRAGMA foreign_key_list(\"{}\")", table.replace('"', "\"\""));
pool.with_connection(|conn| {
let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?;
let rows = stmt
.query_map([], |row| {
Ok(ForeignKeyInfo {
name: format!("fk_{}", row.get::<_, i32>("id")?),
column: row.get("from")?,
ref_schema: None,
ref_table: row.get("table")?,
ref_column: row.get("to")?,
on_update: None,
on_delete: None,
})
})
.map_err(|e| e.to_string())?;
rows.collect::<Result<Vec<_>, _>>().map_err(|e| e.to_string())
})
})
.await
.map_err(|e| e.to_string())?
}
pub async fn list_triggers(pool: &SqliteHandle, _schema: &str, table: &str) -> Result<Vec<TriggerInfo>, String> {
let pool = pool.clone();
let table = table.to_string();
tokio::task::spawn_blocking(move || {
pool.with_connection(|conn| {
let mut stmt = conn
.prepare("SELECT name, sql FROM sqlite_master WHERE type = 'trigger' AND tbl_name = ? ORDER BY name")
.map_err(|e| e.to_string())?;
let rows = stmt
.query_map([table], |row| {
let sql_text: Option<String> = row.get("sql")?;
let upper = sql_text.clone().unwrap_or_default().to_uppercase();
let timing = if upper.contains("BEFORE") {
"BEFORE"
} else if upper.contains("AFTER") {
"AFTER"
} else {
"INSTEAD OF"
};
let event = if upper.contains("INSERT") {
"INSERT"
} else if upper.contains("UPDATE") {
"UPDATE"
} else {
"DELETE"
};
Ok(TriggerInfo {
name: row.get("name")?,
event: event.to_string(),
timing: timing.to_string(),
statement: sql_text,
})
})
.map_err(|e| e.to_string())?;
rows.collect::<Result<Vec<_>, _>>().map_err(|e| e.to_string())
})
})
.await
.map_err(|e| e.to_string())?
}
pub async fn execute_query(pool: &SqliteHandle, sql: &str) -> Result<QueryResult, String> {
execute_query_with_max_rows(pool, sql, None).await
}
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(crate::query::MAX_ROWS).max(1)
}
const SQLITE_FUNCTION_ALIASES: &[(&str, &str)] = &[("if", "IIF"), ("substring", "substr")];
fn normalize_sqlite_sql(sql: &str) -> String {
let mut result = String::with_capacity(sql.len());
let chars: Vec<char> = sql.chars().collect();
let len = chars.len();
let mut i = 0;
while i < len {
if i + 1 < len && chars[i] == '-' && chars[i + 1] == '-' {
while i < len && chars[i] != '\n' {
result.push(chars[i]);
i += 1;
}
continue;
}
if i + 1 < len && chars[i] == '/' && chars[i + 1] == '*' {
while i + 1 < len && !(chars[i] == '*' && chars[i + 1] == '/') {
result.push(chars[i]);
i += 1;
}
if i + 1 < len {
result.push(chars[i]);
result.push(chars[i + 1]);
i += 2;
}
continue;
}
if chars[i] == '\'' {
result.push(chars[i]);
i += 1;
while i < len {
if chars[i] == '\'' {
result.push('\'');
i += 1;
if i < len && chars[i] == '\'' {
result.push('\'');
i += 1;
} else {
break;
}
} else {
result.push(chars[i]);
i += 1;
}
}
continue;
}
let prev = if i == 0 { '\0' } else { chars[i - 1] };
let boundary = !prev.is_alphanumeric() && prev != '_' && prev != '.';
if boundary {
let remaining: String = chars[i..].iter().collect();
let remaining_lower = remaining.to_lowercase();
let mut matched = false;
for (source, replacement) in SQLITE_FUNCTION_ALIASES {
if remaining_lower.starts_with(*source) && chars.get(i + source.len()) != Some(&'_') {
let mut j = i + source.len();
while j < len && chars[j].is_whitespace() {
j += 1;
}
if j < len && chars[j] == '(' {
let whitespace: String = chars[i + source.len()..j].iter().collect();
result.push_str(replacement);
result.push_str(&whitespace);
i = j;
matched = true;
break;
}
}
}
if matched {
continue;
}
}
result.push(chars[i]);
i += 1;
}
result
}
pub async fn execute_query_with_max_rows(
pool: &SqliteHandle,
sql: &str,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let pool = pool.clone();
let sql = normalize_sqlite_sql(sql);
tokio::task::spawn_blocking(move || execute_query_blocking(&pool, &sql, max_rows))
.await
.map_err(|e| e.to_string())?
}
fn execute_query_blocking(pool: &SqliteHandle, sql: &str, max_rows: Option<usize>) -> Result<QueryResult, String> {
let start = Instant::now();
let row_limit = query_result_row_limit(max_rows);
pool.with_connection(|conn| {
if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) {
let mut stmt = conn.prepare(sql).map_err(|e| e.to_string())?;
let columns = stmt.column_names().iter().map(|name| name.to_string()).collect::<Vec<_>>();
let mut rows = stmt.query([]).map_err(|e| e.to_string())?;
let mut result_rows = Vec::new();
while let Some(row) = rows.next().map_err(|e| e.to_string())? {
let mut values = Vec::with_capacity(columns.len());
for i in 0..columns.len() {
values.push(value_ref_to_json(row.get_ref(i).map_err(|e| e.to_string())?));
}
result_rows.push(values);
if result_rows.len() > row_limit {
break;
}
}
let truncated = result_rows.len() > row_limit;
if truncated {
result_rows.truncate(row_limit);
}
Ok(QueryResult {
columns,
column_types: Vec::new(),
column_sortables: vec![],
rows: result_rows,
affected_rows: 0,
execution_time_ms: start.elapsed().as_millis(),
truncated,
session_id: None,
has_more: false,
})
} else {
conn.execute_batch(sql).map_err(|e| e.to_string())?;
Ok(QueryResult {
columns: vec![],
column_types: Vec::new(),
column_sortables: vec![],
rows: vec![],
affected_rows: conn.changes(),
execution_time_ms: start.elapsed().as_millis(),
truncated: false,
session_id: None,
has_more: false,
})
}
})
}
fn value_ref_to_json(value: ValueRef<'_>) -> serde_json::Value {
match value {
ValueRef::Null => serde_json::Value::Null,
ValueRef::Integer(v) => super::safe_i64_to_json(v),
ValueRef::Real(v) => {
serde_json::Number::from_f64(v).map(serde_json::Value::Number).unwrap_or(serde_json::Value::Null)
}
ValueRef::Text(v) => serde_json::Value::String(String::from_utf8_lossy(v).to_string()),
ValueRef::Blob(v) => super::binary_value_to_json(v),
}
}