use duckdb::Connection;
const MARGIN: usize = 200;
const FIND_CAP: usize = 10_000;
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
enum SourceKind {
Parquet,
Json,
Sqlite,
DuckDb,
}
struct Sheet {
name: String,
relation: String,
}
struct WindowCache {
start: usize,
rows: Vec<Vec<String>>,
}
enum Engine {
Duck(DuckBook),
Sqlite(SqliteBook),
}
pub struct DataBook {
engine: Engine,
}
impl DataBook {
pub fn open(path: &str) -> Result<DataBook, String> {
let kind =
detect_kind(path).ok_or_else(|| format!("unrecognised data-file extension: {path}"))?;
let engine = match kind {
SourceKind::Sqlite => Engine::Sqlite(SqliteBook::open(path)?),
_ => Engine::Duck(DuckBook::open(path, kind)?),
};
Ok(DataBook { engine })
}
pub fn names(&self) -> Vec<String> {
match &self.engine {
Engine::Duck(b) => b.names(),
Engine::Sqlite(b) => b.names(),
}
}
pub fn selected(&self) -> usize {
match &self.engine {
Engine::Duck(b) => b.selected(),
Engine::Sqlite(b) => b.selected(),
}
}
pub fn select(&mut self, idx: usize) {
match &mut self.engine {
Engine::Duck(b) => b.select(idx),
Engine::Sqlite(b) => b.select(idx),
}
}
pub fn set_sql(&mut self, sql: &str) -> Result<(), String> {
match &mut self.engine {
Engine::Duck(b) => b.set_sql(sql),
Engine::Sqlite(b) => b.set_sql(sql),
}
}
pub fn active_sql(&self) -> Option<&str> {
match &self.engine {
Engine::Duck(b) => b.active_sql(),
Engine::Sqlite(b) => b.active_sql(),
}
}
pub fn dims(&self) -> (usize, usize, bool, bool) {
match &self.engine {
Engine::Duck(b) => b.dims(),
Engine::Sqlite(b) => b.dims(),
}
}
pub fn headers(&self) -> Vec<String> {
match &self.engine {
Engine::Duck(b) => b.headers(),
Engine::Sqlite(b) => b.headers(),
}
}
pub fn window(&mut self, r0: usize, r1: usize, c0: usize, c1: usize) -> Vec<Vec<String>> {
match &mut self.engine {
Engine::Duck(b) => b.window(r0, r1, c0, c1),
Engine::Sqlite(b) => b.window(r0, r1, c0, c1),
}
}
pub fn find(&self, query: &str) -> Vec<(usize, usize)> {
match &self.engine {
Engine::Duck(b) => b.find(query),
Engine::Sqlite(b) => b.find(query),
}
}
}
struct DuckBook {
conn: Connection,
sheets: Vec<Sheet>,
cur: usize,
override_sql: Option<String>,
schema: Vec<(String, String)>,
total: usize,
cache: Option<WindowCache>,
}
impl DuckBook {
fn open(path: &str, kind: SourceKind) -> Result<DuckBook, String> {
let conn = open_conn()?;
let sheets = register(&conn, kind, path)?;
if sheets.is_empty() {
return Err("database has no tables".into());
}
let mut book = DuckBook {
conn,
sheets,
cur: 0,
override_sql: None,
schema: Vec::new(),
total: 0,
cache: None,
};
book.load_active()?;
Ok(book)
}
fn active_relation(&self) -> String {
match &self.override_sql {
Some(sql) => sql.clone(),
None => self.sheets[self.cur].relation.clone(),
}
}
fn load_active(&mut self) -> Result<(), String> {
let relation = self.active_relation();
self.schema = describe(&self.conn, &relation)?;
self.total = count(&self.conn, &relation)?;
self.cache = None;
Ok(())
}
fn names(&self) -> Vec<String> {
self.sheets.iter().map(|s| s.name.clone()).collect()
}
fn selected(&self) -> usize {
self.cur
}
fn select(&mut self, idx: usize) {
if idx < self.sheets.len() && idx != self.cur {
self.cur = idx;
self.override_sql = None;
let _ = self.load_active();
}
}
fn set_sql(&mut self, sql: &str) -> Result<(), String> {
let trimmed = sql.trim();
if trimmed.is_empty() {
self.override_sql = None;
self.load_active()?;
return Ok(());
}
let prev = self.override_sql.take();
self.override_sql = Some(trimmed.to_string());
match self.load_active() {
Ok(()) => Ok(()),
Err(e) => {
self.override_sql = prev;
let _ = self.load_active();
Err(e)
}
}
}
fn active_sql(&self) -> Option<&str> {
self.override_sql.as_deref()
}
fn dims(&self) -> (usize, usize, bool, bool) {
(self.total, self.schema.len(), true, false)
}
fn headers(&self) -> Vec<String> {
self.schema.iter().map(|(n, _)| n.clone()).collect()
}
fn window(&mut self, r0: usize, r1: usize, c0: usize, c1: usize) -> Vec<Vec<String>> {
let relation = self.active_relation();
let conn = &self.conn;
let schema = &self.schema;
window_cached(&mut self.cache, self.total, r0, r1, c0, c1, |start, len| {
fetch_duck(conn, schema, &relation, start, len)
})
}
fn find(&self, query: &str) -> Vec<(usize, usize)> {
let ncols = self.schema.len();
if ncols == 0 || query.is_empty() {
return Vec::new();
}
let needle = query.to_ascii_lowercase();
let sql = match self.find_sql(query) {
Some(s) => s,
None => return Vec::new(),
};
let mut hits = Vec::new();
let Ok(mut stmt) = self.conn.prepare(&sql) else {
return hits;
};
let Ok(mut rows) = stmt.query([]) else {
return hits;
};
while let Ok(Some(row)) = rows.next() {
let Ok(rn) = row.get::<usize, i64>(0) else {
continue;
};
let r = rn.max(0) as usize;
for c in 0..ncols {
let cell: Option<String> = row.get(c + 1).unwrap_or(None);
if let Some(text) = cell {
if crate::xlsx::contains_ci(&text, &needle) {
hits.push((r, c));
if hits.len() >= FIND_CAP {
return hits;
}
}
}
}
}
hits
}
fn find_sql(&self, query: &str) -> Option<String> {
if self.schema.is_empty() {
return None;
}
let rel = self.active_relation();
let pattern = quote_literal(&format!("%{}%", escape_like(query)));
let mut proj = String::from("__sucher_rn");
let mut filter = String::new();
for (i, (name, _)) in self.schema.iter().enumerate() {
let cast = format!("CAST({} AS VARCHAR)", quote_ident(name));
proj.push_str(", ");
proj.push_str(&cast);
if i > 0 {
filter.push_str(" OR ");
}
filter.push_str(&format!("({cast} ILIKE {pattern} ESCAPE '\\')"));
}
Some(format!(
"SELECT {proj} FROM \
(SELECT (row_number() OVER () - 1) AS __sucher_rn, * FROM ({rel})) \
WHERE {filter} LIMIT {FIND_CAP}"
))
}
}
fn fetch_duck(
conn: &Connection,
schema: &[(String, String)],
relation: &str,
start: usize,
len: usize,
) -> Result<Vec<Vec<String>>, String> {
let ncols = schema.len();
if ncols == 0 {
return Ok(Vec::new());
}
let proj = cast_projection(schema);
let sql = format!("SELECT {proj} FROM ({relation}) LIMIT {len} OFFSET {start}");
let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?;
let mut rows = stmt.query([]).map_err(|e| e.to_string())?;
let mut out = Vec::new();
while let Some(row) = rows.next().map_err(|e| e.to_string())? {
let mut line = Vec::with_capacity(ncols);
for c in 0..ncols {
let v: Option<String> = row.get(c).map_err(|e| e.to_string())?;
line.push(v.unwrap_or_default());
}
out.push(line);
}
Ok(out)
}
fn open_conn() -> Result<Connection, String> {
let conn = Connection::open_in_memory().map_err(|e| e.to_string())?;
conn.execute_batch(
"SET autoinstall_known_extensions=false; SET autoload_known_extensions=false;",
)
.map_err(|e| e.to_string())?;
Ok(conn)
}
fn register(conn: &Connection, kind: SourceKind, path: &str) -> Result<Vec<Sheet>, String> {
let lit = quote_literal(path);
match kind {
SourceKind::Parquet => Ok(vec![single_view(
conn,
path,
&format!("read_parquet({lit})"),
)?]),
SourceKind::Json => Ok(vec![single_view(
conn,
path,
&format!("read_json_auto({lit})"),
)?]),
SourceKind::DuckDb => {
conn.execute_batch(&format!("ATTACH {lit} AS db (READ_ONLY);"))
.map_err(|e| e.to_string())?;
attached_sheets(conn)
}
SourceKind::Sqlite => {
unreachable!("SQLite is read by SqliteBook, never routed to DuckBook")
}
}
}
fn single_view(conn: &Connection, path: &str, reader: &str) -> Result<Sheet, String> {
let stem = file_stem(path);
let ident = quote_ident(&stem);
conn.execute_batch(&format!("CREATE VIEW {ident} AS SELECT * FROM {reader};"))
.map_err(|e| e.to_string())?;
Ok(Sheet {
name: stem,
relation: format!("SELECT * FROM {ident}"),
})
}
fn attached_sheets(conn: &Connection) -> Result<Vec<Sheet>, String> {
let mut stmt = conn
.prepare(
"SELECT table_name FROM information_schema.tables \
WHERE table_catalog='db' ORDER BY 1",
)
.map_err(|e| e.to_string())?;
let mut rows = stmt.query([]).map_err(|e| e.to_string())?;
let mut sheets = Vec::new();
while let Some(row) = rows.next().map_err(|e| e.to_string())? {
let name: String = row.get(0).map_err(|e| e.to_string())?;
let relation = format!("SELECT * FROM db.{}", quote_ident(&name));
sheets.push(Sheet { name, relation });
}
Ok(sheets)
}
fn describe(conn: &Connection, relation: &str) -> Result<Vec<(String, String)>, String> {
let mut stmt = conn
.prepare(&format!("DESCRIBE {relation}"))
.map_err(|e| e.to_string())?;
let mut rows = stmt.query([]).map_err(|e| e.to_string())?;
let mut schema = Vec::new();
while let Some(row) = rows.next().map_err(|e| e.to_string())? {
let name: String = row.get(0).map_err(|e| e.to_string())?;
let ty: String = row.get(1).map_err(|e| e.to_string())?;
schema.push((name, ty));
}
Ok(schema)
}
fn count(conn: &Connection, relation: &str) -> Result<usize, String> {
let n: i64 = conn
.query_row(&format!("SELECT count(*) FROM ({relation})"), [], |r| {
r.get(0)
})
.map_err(|e| e.to_string())?;
Ok(n.max(0) as usize)
}
struct SqliteBook {
conn: rusqlite::Connection,
sheets: Vec<Sheet>,
cur: usize,
override_sql: Option<String>,
columns: Vec<String>,
total: usize,
cache: Option<WindowCache>,
}
impl SqliteBook {
fn open(path: &str) -> Result<SqliteBook, String> {
let conn =
rusqlite::Connection::open_with_flags(path, rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY)
.map_err(|e| e.to_string())?;
let sheets = sqlite_sheets(&conn)?;
if sheets.is_empty() {
return Err("database has no tables".into());
}
let mut book = SqliteBook {
conn,
sheets,
cur: 0,
override_sql: None,
columns: Vec::new(),
total: 0,
cache: None,
};
book.load_active()?;
Ok(book)
}
fn active_relation(&self) -> String {
match &self.override_sql {
Some(sql) => sql.clone(),
None => self.sheets[self.cur].relation.clone(),
}
}
fn load_active(&mut self) -> Result<(), String> {
let relation = self.active_relation();
self.columns = sqlite_columns(&self.conn, &relation)?;
self.total = sqlite_count(&self.conn, &relation)?;
self.cache = None;
Ok(())
}
fn names(&self) -> Vec<String> {
self.sheets.iter().map(|s| s.name.clone()).collect()
}
fn selected(&self) -> usize {
self.cur
}
fn select(&mut self, idx: usize) {
if idx < self.sheets.len() && idx != self.cur {
self.cur = idx;
self.override_sql = None;
let _ = self.load_active();
}
}
fn set_sql(&mut self, sql: &str) -> Result<(), String> {
let trimmed = sql.trim();
if trimmed.is_empty() {
self.override_sql = None;
self.load_active()?;
return Ok(());
}
let prev = self.override_sql.take();
self.override_sql = Some(trimmed.to_string());
match self.load_active() {
Ok(()) => Ok(()),
Err(e) => {
self.override_sql = prev;
let _ = self.load_active();
Err(e)
}
}
}
fn active_sql(&self) -> Option<&str> {
self.override_sql.as_deref()
}
fn dims(&self) -> (usize, usize, bool, bool) {
(self.total, self.columns.len(), true, false)
}
fn headers(&self) -> Vec<String> {
self.columns.clone()
}
fn window(&mut self, r0: usize, r1: usize, c0: usize, c1: usize) -> Vec<Vec<String>> {
let relation = self.active_relation();
let conn = &self.conn;
let ncols = self.columns.len();
window_cached(&mut self.cache, self.total, r0, r1, c0, c1, |start, len| {
fetch_sqlite(conn, ncols, &relation, start, len)
})
}
fn find(&self, query: &str) -> Vec<(usize, usize)> {
let ncols = self.columns.len();
if ncols == 0 || query.is_empty() {
return Vec::new();
}
let needle = query.to_ascii_lowercase();
let sql = match self.find_sql(query) {
Some(s) => s,
None => return Vec::new(),
};
let mut hits = Vec::new();
let Ok(mut stmt) = self.conn.prepare(&sql) else {
return hits;
};
let Ok(mut rows) = stmt.query([]) else {
return hits;
};
while let Ok(Some(row)) = rows.next() {
let Ok(rn) = row.get::<usize, i64>(0) else {
continue;
};
let r = rn.max(0) as usize;
for c in 0..ncols {
let Ok(cell) = row.get_ref(c + 1) else {
continue;
};
let text = fmt_value_ref(cell);
if crate::xlsx::contains_ci(&text, &needle) {
hits.push((r, c));
if hits.len() >= FIND_CAP {
return hits;
}
}
}
}
hits
}
fn find_sql(&self, query: &str) -> Option<String> {
if self.columns.is_empty() {
return None;
}
let rel = self.active_relation();
let pattern = quote_literal(&format!("%{}%", escape_like(query)));
let mut proj = String::from("__sucher_rn");
let mut filter = String::new();
for (i, name) in self.columns.iter().enumerate() {
let cast = format!("CAST({} AS TEXT)", quote_ident(name));
proj.push_str(", ");
proj.push_str(&cast);
if i > 0 {
filter.push_str(" OR ");
}
filter.push_str(&format!("({cast} LIKE {pattern} ESCAPE '\\')"));
}
Some(format!(
"SELECT {proj} FROM \
(SELECT (row_number() OVER () - 1) AS __sucher_rn, * FROM ({rel})) \
WHERE {filter} LIMIT {FIND_CAP}"
))
}
}
fn sqlite_sheets(conn: &rusqlite::Connection) -> Result<Vec<Sheet>, String> {
let mut stmt = conn
.prepare(
"SELECT name FROM sqlite_master \
WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY 1",
)
.map_err(|e| e.to_string())?;
let names = stmt
.query_map([], |row| row.get::<usize, String>(0))
.map_err(|e| e.to_string())?;
let mut sheets = Vec::new();
for name in names {
let name = name.map_err(|e| e.to_string())?;
let relation = format!("SELECT * FROM {}", quote_ident(&name));
sheets.push(Sheet { name, relation });
}
Ok(sheets)
}
fn sqlite_columns(conn: &rusqlite::Connection, relation: &str) -> Result<Vec<String>, String> {
let stmt = conn
.prepare(&format!("SELECT * FROM ({relation}) LIMIT 0"))
.map_err(|e| e.to_string())?;
Ok(stmt
.column_names()
.into_iter()
.map(|s| s.to_string())
.collect())
}
fn sqlite_count(conn: &rusqlite::Connection, relation: &str) -> Result<usize, String> {
let n: i64 = conn
.query_row(&format!("SELECT count(*) FROM ({relation})"), [], |r| {
r.get(0)
})
.map_err(|e| e.to_string())?;
Ok(n.max(0) as usize)
}
fn fetch_sqlite(
conn: &rusqlite::Connection,
ncols: usize,
relation: &str,
start: usize,
len: usize,
) -> Result<Vec<Vec<String>>, String> {
if ncols == 0 {
return Ok(Vec::new());
}
let sql = format!("SELECT * FROM ({relation}) LIMIT {len} OFFSET {start}");
let mut stmt = conn.prepare(&sql).map_err(|e| e.to_string())?;
let mut rows = stmt.query([]).map_err(|e| e.to_string())?;
let mut out = Vec::new();
while let Some(row) = rows.next().map_err(|e| e.to_string())? {
let mut line = Vec::with_capacity(ncols);
for i in 0..ncols {
let vr = row.get_ref(i).map_err(|e| e.to_string())?;
line.push(fmt_value_ref(vr));
}
out.push(line);
}
Ok(out)
}
fn fmt_value_ref(v: rusqlite::types::ValueRef<'_>) -> String {
use rusqlite::types::ValueRef;
match v {
ValueRef::Null => String::new(),
ValueRef::Integer(i) => i.to_string(),
ValueRef::Real(f) => fmt_real(f),
ValueRef::Text(b) => String::from_utf8_lossy(b).into_owned(),
ValueRef::Blob(b) => format!("[{} bytes]", b.len()),
}
}
fn fmt_real(f: f64) -> String {
if f.fract() == 0.0 && f.abs() < 1e15 {
format!("{}", f as i64)
} else {
format!("{f}")
}
}
fn cache_covers(cache: &Option<WindowCache>, r0: usize, r1: usize) -> bool {
matches!(cache, Some(c) if c.start <= r0 && r1 <= c.start + c.rows.len())
}
fn window_cached(
cache: &mut Option<WindowCache>,
total: usize,
r0: usize,
r1: usize,
c0: usize,
c1: usize,
fetch: impl FnOnce(usize, usize) -> Result<Vec<Vec<String>>, String>,
) -> Vec<Vec<String>> {
let r1 = r1.min(total);
if r0 >= r1 {
return Vec::new();
}
if !cache_covers(cache, r0, r1) {
let start = r0.saturating_sub(MARGIN);
let len = ((r1 - r0) + 2 * MARGIN).min(total.saturating_sub(start));
match fetch(start, len) {
Ok(rows) => *cache = Some(WindowCache { start, rows }),
Err(_) => {
*cache = None;
return Vec::new();
}
}
}
let cache = cache.as_ref().expect("cache populated above");
(r0..r1)
.map(|r| {
let row = cache.rows.get(r - cache.start);
(c0..c1)
.map(|c| {
row.and_then(|cells| cells.get(c))
.cloned()
.unwrap_or_default()
})
.collect()
})
.collect()
}
fn detect_kind(path: &str) -> Option<SourceKind> {
let ext = std::path::Path::new(path)
.extension()?
.to_str()?
.to_ascii_lowercase();
match ext.as_str() {
"parquet" | "pq" => Some(SourceKind::Parquet),
"jsonl" | "ndjson" => Some(SourceKind::Json),
"sqlite" | "sqlite3" | "db" | "db3" => Some(SourceKind::Sqlite),
"duckdb" | "ddb" => Some(SourceKind::DuckDb),
_ => None,
}
}
fn file_stem(path: &str) -> String {
std::path::Path::new(path)
.file_stem()
.map(|s| s.to_string_lossy().into_owned())
.filter(|s| !s.is_empty())
.unwrap_or_else(|| "data".to_string())
}
fn quote_ident(s: &str) -> String {
format!("\"{}\"", s.replace('"', "\"\""))
}
fn quote_literal(s: &str) -> String {
format!("'{}'", s.replace('\'', "''"))
}
fn cast_projection(schema: &[(String, String)]) -> String {
schema
.iter()
.map(|(name, _)| format!("CAST({} AS VARCHAR)", quote_ident(name)))
.collect::<Vec<_>>()
.join(", ")
}
fn escape_like(s: &str) -> String {
let mut out = String::with_capacity(s.len());
for c in s.chars() {
match c {
'\\' => out.push_str("\\\\"),
'%' => out.push_str("\\%"),
'_' => out.push_str("\\_"),
_ => out.push(c),
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn quote_ident_doubles_quotes() {
assert_eq!(quote_ident("plain"), "\"plain\"");
assert_eq!(quote_ident("we\"ird"), "\"we\"\"ird\"");
assert_eq!(quote_ident("select.from"), "\"select.from\"");
}
#[test]
fn quote_literal_doubles_apostrophes() {
assert_eq!(quote_literal("/tmp/a.parquet"), "'/tmp/a.parquet'");
assert_eq!(quote_literal("it's"), "'it''s'");
}
#[test]
fn cast_projection_builds_varchar_list() {
let schema = vec![
("id".to_string(), "INTEGER".to_string()),
("na\"me".to_string(), "VARCHAR".to_string()),
];
assert_eq!(
cast_projection(&schema),
"CAST(\"id\" AS VARCHAR), CAST(\"na\"\"me\" AS VARCHAR)"
);
assert_eq!(cast_projection(&[]), "");
}
#[test]
fn detect_kind_from_extension() {
assert_eq!(detect_kind("a.parquet"), Some(SourceKind::Parquet));
assert_eq!(detect_kind("a.PQ"), Some(SourceKind::Parquet));
assert_eq!(detect_kind("events.jsonl"), Some(SourceKind::Json));
assert_eq!(detect_kind("events.ndjson"), Some(SourceKind::Json));
assert_eq!(detect_kind("app.db"), Some(SourceKind::Sqlite));
assert_eq!(detect_kind("app.sqlite3"), Some(SourceKind::Sqlite));
assert_eq!(detect_kind("store.duckdb"), Some(SourceKind::DuckDb));
assert_eq!(detect_kind("store.ddb"), Some(SourceKind::DuckDb));
assert_eq!(detect_kind("notes.txt"), None);
assert_eq!(detect_kind("noext"), None);
}
#[test]
fn file_stem_falls_back() {
assert_eq!(file_stem("/data/sales.parquet"), "sales");
assert_eq!(file_stem("events.jsonl"), "events");
assert_eq!(file_stem("/x/.parquet"), ".parquet");
assert_eq!(file_stem(""), "data");
assert_eq!(file_stem("/"), "data");
}
#[test]
fn escape_like_escapes_wildcards() {
assert_eq!(escape_like("plain"), "plain");
assert_eq!(escape_like("50%"), "50\\%");
assert_eq!(escape_like("a_b"), "a\\_b");
assert_eq!(escape_like("back\\slash"), "back\\\\slash");
}
#[test]
fn fmt_value_ref_formats_each_type() {
use rusqlite::types::ValueRef;
assert_eq!(fmt_value_ref(ValueRef::Null), "");
assert_eq!(fmt_value_ref(ValueRef::Integer(42)), "42");
assert_eq!(fmt_value_ref(ValueRef::Real(1.5)), "1.5");
assert_eq!(fmt_value_ref(ValueRef::Real(3.0)), "3");
assert_eq!(fmt_value_ref(ValueRef::Text(b"h\xC3\xA9llo")), "héllo");
assert_eq!(fmt_value_ref(ValueRef::Text(b"a\xFFb")), "a\u{FFFD}b");
assert_eq!(fmt_value_ref(ValueRef::Blob(&[1, 2, 3, 4])), "[4 bytes]");
}
fn scratch_path(tag: &str, ext: &str) -> std::path::PathBuf {
let mut p = std::env::temp_dir();
p.push(format!(
"sucher-data-{tag}-{}-{}.{ext}",
std::process::id(),
COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
));
p
}
static COUNTER: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
fn make_parquet() -> std::path::PathBuf {
let path = scratch_path("pq", "parquet");
let conn = open_conn().unwrap();
conn.execute_batch(&format!(
"COPY (SELECT * FROM (VALUES \
(1, 'alpha', DATE '2020-01-02', 'x'), \
(2, 'beta', DATE '2021-06-15', NULL), \
(3, 'gamma', DATE '2022-12-31', 'z')) \
t(id, name, d, note)) TO {} (FORMAT PARQUET);",
quote_literal(path.to_str().unwrap())
))
.unwrap();
path
}
fn make_sqlite() -> std::path::PathBuf {
let path = scratch_path("sq", "db");
{
let conn = rusqlite::Connection::open(&path).unwrap();
conn.execute_batch(
"CREATE TABLE items (id INTEGER PRIMARY KEY AUTOINCREMENT, \
label TEXT, note TEXT); \
INSERT INTO items (id, label, note) VALUES (1, 'apple', 'x'); \
INSERT INTO items (id, label, note) VALUES (2, 'pear', NULL); \
CREATE TABLE tags (name TEXT); \
INSERT INTO tags VALUES ('fresh');",
)
.unwrap();
} path
}
#[test]
fn parquet_dims_headers_window_find() {
let path = make_parquet();
let p = path.to_str().unwrap();
let mut book = DataBook::open(p).expect("open parquet");
assert_eq!(book.names().len(), 1);
let (total, ncols, done, capped) = book.dims();
assert_eq!(total, 3);
assert_eq!(ncols, 4);
assert!(done);
assert!(!capped);
assert_eq!(book.headers(), vec!["id", "name", "d", "note"]);
let win = book.window(0, 3, 0, 4);
assert_eq!(win[0], vec!["1", "alpha", "2020-01-02", "x"]);
assert_eq!(win[1], vec!["2", "beta", "2021-06-15", ""]); assert_eq!(win[2][2], "2022-12-31");
let hits = book.find("BETA");
assert_eq!(hits, vec![(1, 1)]);
assert!(book.find("zzznotfoundzzz").is_empty());
std::fs::remove_file(path).ok();
}
#[test]
fn set_sql_overrides_view_and_reverts() {
let path = make_parquet();
let p = path.to_str().unwrap();
let mut book = DataBook::open(p).expect("open parquet");
let view = quote_ident(&file_stem(p));
assert_eq!(book.dims(), (3, 4, true, false));
assert_eq!(book.active_sql(), None);
let q = format!("SELECT count(*) AS n FROM {view}");
book.set_sql(&q).expect("ok");
assert_eq!(book.dims(), (1, 1, true, false));
assert_eq!(book.headers(), vec!["n"]);
assert_eq!(book.window(0, 1, 0, 1), vec![vec!["3"]]);
assert_eq!(book.active_sql(), Some(q.as_str()));
assert_eq!(book.find("3"), vec![(0, 0)]);
book.set_sql("").expect("revert ok");
assert_eq!(book.dims(), (3, 4, true, false));
assert_eq!(book.headers(), vec!["id", "name", "d", "note"]);
assert_eq!(book.active_sql(), None);
std::fs::remove_file(path).ok();
}
#[test]
fn set_sql_error_restores_previous_view() {
let path = make_parquet();
let p = path.to_str().unwrap();
let mut book = DataBook::open(p).expect("open parquet");
let view = quote_ident(&file_stem(p));
let good = format!("SELECT id, name FROM {view} WHERE id >= 2");
book.set_sql(&good).expect("ok");
assert_eq!(book.dims(), (2, 2, true, false));
assert_eq!(book.headers(), vec!["id", "name"]);
let before = book.window(0, 2, 0, 2);
let err = book
.set_sql(&format!("SELECT nonexistent_col FROM {view}"))
.expect_err("bad query errors");
assert!(!err.is_empty(), "a human-readable DuckDB error is returned");
assert_eq!(book.dims(), (2, 2, true, false));
assert_eq!(book.headers(), vec!["id", "name"]);
assert_eq!(book.window(0, 2, 0, 2), before);
assert_eq!(book.active_sql(), Some(good.as_str()));
std::fs::remove_file(path).ok();
}
#[test]
fn sqlite_tables_become_sheets() {
let path = make_sqlite();
let p = path.to_str().unwrap();
let mut book = DataBook::open(p).expect("open sqlite");
assert_eq!(book.names(), vec!["items", "tags"]);
book.select(0);
assert_eq!(book.dims(), (2, 3, true, false));
assert_eq!(book.headers(), vec!["id", "label", "note"]);
let win = book.window(0, 2, 0, 3);
assert_eq!(win[0], vec!["1", "apple", "x"]);
assert_eq!(win[1], vec!["2", "pear", ""]);
assert_eq!(book.find("APPLE"), vec![(0, 1)]);
assert!(book.find("zzznope").is_empty());
std::fs::remove_file(path).ok();
}
#[test]
fn sqlite_set_sql_override_revert_and_restore() {
let path = make_sqlite();
let p = path.to_str().unwrap();
let mut book = DataBook::open(p).expect("open sqlite");
assert_eq!(book.dims(), (2, 3, true, false));
assert_eq!(book.active_sql(), None);
let q = "SELECT count(*) AS n FROM items";
book.set_sql(q).expect("ok");
assert_eq!(book.dims(), (1, 1, true, false));
assert_eq!(book.headers(), vec!["n"]);
assert_eq!(book.window(0, 1, 0, 1), vec![vec!["2"]]);
assert_eq!(book.active_sql(), Some(q));
book.set_sql("").expect("revert ok");
assert_eq!(book.dims(), (2, 3, true, false));
assert_eq!(book.headers(), vec!["id", "label", "note"]);
assert_eq!(book.active_sql(), None);
let good = "SELECT id, label FROM items WHERE id >= 2";
book.set_sql(good).expect("ok");
assert_eq!(book.dims(), (1, 2, true, false));
let before = book.window(0, 1, 0, 2);
let err = book
.set_sql("SELECT nonexistent_col FROM items")
.expect_err("bad query errors");
assert!(!err.is_empty(), "a human-readable SQLite error is returned");
assert_eq!(book.dims(), (1, 2, true, false));
assert_eq!(book.headers(), vec!["id", "label"]);
assert_eq!(book.window(0, 1, 0, 2), before);
assert_eq!(book.active_sql(), Some(good));
std::fs::remove_file(path).ok();
}
#[test]
fn sqlite_select_clears_override() {
let path = make_sqlite();
let p = path.to_str().unwrap();
let mut book = DataBook::open(p).expect("open sqlite");
book.set_sql("SELECT count(*) AS n FROM items").expect("ok");
assert_eq!(book.active_sql(), Some("SELECT count(*) AS n FROM items"));
book.select(1); assert_eq!(book.active_sql(), None);
assert_eq!(book.dims(), (1, 1, true, false));
assert_eq!(book.headers(), vec!["name"]);
assert_eq!(book.window(0, 1, 0, 1), vec![vec!["fresh"]]);
std::fs::remove_file(path).ok();
}
#[test]
fn window_prefetch_cache_serves_scroll() {
let path = scratch_path("big", "parquet");
let p = path.to_str().unwrap();
let conn = open_conn().unwrap();
conn.execute_batch(&format!(
"COPY (SELECT i AS id, ('row' || i) AS name FROM range(5000) t(i)) \
TO {} (FORMAT PARQUET);",
quote_literal(p)
))
.unwrap();
drop(conn);
let mut book = DataBook::open(p).expect("open");
assert_eq!(book.dims().0, 5000);
let a = book.window(100, 110, 0, 2);
assert_eq!(a[0], vec!["100", "row100"]);
let b = book.window(101, 111, 0, 2);
assert_eq!(b[0], vec!["101", "row101"]);
let c = book.window(4990, 5000, 0, 2);
assert_eq!(
c.last().unwrap(),
&vec!["4999".to_string(), "row4999".to_string()]
);
std::fs::remove_file(path).ok();
}
#[test]
#[ignore = "set SUCHER_BIG_PARQUET to a large .parquet to run"]
fn big_parquet_opens_and_scrolls_lazily() {
let Ok(path) = std::env::var("SUCHER_BIG_PARQUET") else {
return;
};
let t = std::time::Instant::now();
let mut book = DataBook::open(&path).expect("open big parquet");
let (total, ncols, _, _) = book.dims();
eprintln!(
"open {total} rows × {ncols} cols in {} ms",
t.elapsed().as_millis()
);
let t2 = std::time::Instant::now();
let start = total.saturating_sub(50);
let win = book.window(start, total, 0, ncols.min(8));
eprintln!(
"window @{start} ({} rows) in {} ms",
win.len(),
t2.elapsed().as_millis()
);
assert!(!win.is_empty());
}
#[test]
fn offline_reads_parquet_without_network() {
let path = make_parquet();
let p = path.to_str().unwrap();
let mut extdir = std::env::temp_dir();
extdir.push(format!(
"sucher-empty-extdir-{}-{}",
std::process::id(),
COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed)
));
std::fs::create_dir_all(&extdir).unwrap();
let conn = open_conn().unwrap();
conn.execute_batch(&format!(
"SET extension_directory={};",
quote_literal(extdir.to_str().unwrap())
))
.expect("set empty extension_directory");
let n: i64 = conn
.query_row(
&format!("SELECT count(*) FROM read_parquet({})", quote_literal(p)),
[],
|r| r.get(0),
)
.expect("read_parquet must work offline from the static build");
assert_eq!(n, 3);
std::fs::remove_file(path).ok();
std::fs::remove_dir_all(extdir).ok();
}
}