use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use rusqlite::types::ValueRef;
use rusqlite::Connection;
use crate::error::Result;
use crate::ingest::quote_ident;
use crate::store::{DataStore, IngestReport, QueryResult, TableInfo, TableSummary};
use crate::{ingest, summary};
const DATA_DIR: &str = ".newt-data";
const DB_FILE: &str = "data.db";
#[derive(Debug, Clone)]
pub struct SqliteBackend {
conn: Arc<Mutex<Connection>>,
}
impl SqliteBackend {
pub fn open(path: &Path) -> Result<Self> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let conn = Connection::open(path)?;
apply_journal_mode(&conn)?;
ensure_metadata_table(&conn)?;
Ok(Self {
conn: Arc::new(Mutex::new(conn)),
})
}
pub fn open_in_memory() -> Result<Self> {
let conn = Connection::open_in_memory()?;
ensure_metadata_table(&conn)?;
Ok(Self {
conn: Arc::new(Mutex::new(conn)),
})
}
pub fn default_db_path(workspace: &Path) -> PathBuf {
workspace.join(DATA_DIR).join(DB_FILE)
}
fn lock_conn(&self) -> std::sync::MutexGuard<'_, Connection> {
self.conn
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
impl DataStore for SqliteBackend {
fn ingest_csv(&self, path: &Path, table: &str) -> Result<IngestReport> {
let conn = self.lock_conn();
ingest::ingest_csv(&conn, path, table)
}
fn query(&self, sql: &str, row_cap: usize) -> Result<QueryResult> {
let conn = self.lock_conn();
let mut stmt = conn.prepare(sql)?;
let columns: Vec<String> = stmt.column_names().iter().map(|c| c.to_string()).collect();
if stmt.column_count() == 0 {
stmt.raw_execute()?;
return Ok(QueryResult {
columns,
rows: Vec::new(),
returned: 0,
truncated: false,
});
}
let mut sql_rows = stmt.query([])?;
let mut rows: Vec<Vec<serde_json::Value>> = Vec::new();
let mut truncated = false;
while let Some(row) = sql_rows.next()? {
if rows.len() >= row_cap {
truncated = true;
break;
}
let mut out_row = Vec::with_capacity(columns.len());
for i in 0..columns.len() {
out_row.push(value_ref_to_json(row.get_ref(i)?));
}
rows.push(out_row);
}
let returned = rows.len();
Ok(QueryResult {
columns,
rows,
returned,
truncated,
})
}
fn summarize(&self, table: &str) -> Result<TableSummary> {
let conn = self.lock_conn();
summary::summarize(&conn, table)
}
fn list_tables(&self) -> Result<Vec<TableInfo>> {
let conn = self.lock_conn();
let names: Vec<String> = {
let mut stmt = conn.prepare(
"SELECT name FROM sqlite_master
WHERE type = 'table'
AND name NOT LIKE 'sqlite_%'
AND name != '__datasets'
ORDER BY name",
)?;
let rows = stmt.query_map([], |row| row.get::<_, String>(0))?;
rows.collect::<rusqlite::Result<Vec<_>>>()?
};
let mut infos = Vec::with_capacity(names.len());
for name in names {
let quoted = quote_ident(&name);
let row_count: u64 =
conn.query_row(&format!("SELECT COUNT(*) FROM {quoted}"), [], |row| {
row.get::<_, i64>(0).map(|n| n.max(0) as u64)
})?;
let source: Option<String> = conn
.query_row(
"SELECT source FROM __datasets WHERE table_name = ?1",
[&name],
|row| row.get::<_, Option<String>>(0),
)
.ok()
.flatten();
infos.push(TableInfo {
table: name,
row_count,
source,
});
}
Ok(infos)
}
fn backend_name(&self) -> &'static str {
"sqlite"
}
}
fn apply_journal_mode(conn: &Connection) -> Result<()> {
let wal: rusqlite::Result<String> =
conn.pragma_update_and_check(None, "journal_mode", "WAL", |row| row.get(0));
match wal {
Ok(mode) if mode.eq_ignore_ascii_case("wal") => {
conn.pragma_update(None, "synchronous", "NORMAL")?;
Ok(())
}
Ok(mode) => {
tracing::warn!(%mode, "journal_mode=WAL did not take; keeping synchronous=FULL");
Ok(())
}
Err(e) if wal_fallback_eligible(&e.to_string()) => {
tracing::warn!(
error = %e,
"SQLite refused WAL (network filesystem?); data.db is running on the \
slower journal_mode=DELETE fallback"
);
conn.pragma_update(None, "journal_mode", "DELETE")?;
Ok(())
}
Err(e) => Err(e.into()),
}
}
fn wal_fallback_eligible(error_text: &str) -> bool {
let lower = error_text.to_lowercase();
lower.contains("locking protocol") || lower.contains("disk i/o error")
}
fn ensure_metadata_table(conn: &Connection) -> Result<()> {
conn.execute(
"CREATE TABLE IF NOT EXISTS __datasets (
table_name TEXT PRIMARY KEY,
source TEXT,
row_count INTEGER NOT NULL,
ingested_at INTEGER NOT NULL
)",
[],
)?;
Ok(())
}
fn value_ref_to_json(value: ValueRef<'_>) -> serde_json::Value {
use serde_json::Value;
match value {
ValueRef::Null => Value::Null,
ValueRef::Integer(i) => Value::Number(i.into()),
ValueRef::Real(f) => serde_json::Number::from_f64(f)
.map(Value::Number)
.unwrap_or(Value::Null),
ValueRef::Text(bytes) => Value::String(String::from_utf8_lossy(bytes).into_owned()),
ValueRef::Blob(bytes) => Value::String(base64_encode(bytes)),
}
}
fn base64_encode(bytes: &[u8]) -> String {
const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
let mut out = String::with_capacity(bytes.len().div_ceil(3) * 4);
for chunk in bytes.chunks(3) {
let b0 = chunk[0];
let b1 = chunk.get(1).copied().unwrap_or(0);
let b2 = chunk.get(2).copied().unwrap_or(0);
out.push(ALPHABET[(b0 >> 2) as usize] as char);
out.push(ALPHABET[(((b0 & 0b11) << 4) | (b1 >> 4)) as usize] as char);
if chunk.len() > 1 {
out.push(ALPHABET[(((b1 & 0b1111) << 2) | (b2 >> 6)) as usize] as char);
} else {
out.push('=');
}
if chunk.len() > 2 {
out.push(ALPHABET[(b2 & 0b111111) as usize] as char);
} else {
out.push('=');
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn base64_matches_known_vectors() {
assert_eq!(base64_encode(b""), "");
assert_eq!(base64_encode(b"f"), "Zg==");
assert_eq!(base64_encode(b"fo"), "Zm8=");
assert_eq!(base64_encode(b"foo"), "Zm9v");
assert_eq!(base64_encode(b"foob"), "Zm9vYg==");
assert_eq!(base64_encode(b"fooba"), "Zm9vYmE=");
assert_eq!(base64_encode(b"foobar"), "Zm9vYmFy");
}
#[test]
fn default_db_path_is_under_newt_data() {
let p = SqliteBackend::default_db_path(Path::new("/tmp/ws"));
assert_eq!(p, Path::new("/tmp/ws/.newt-data/data.db"));
}
#[test]
fn wal_fallback_predicate() {
assert!(wal_fallback_eligible("disk I/O error"));
assert!(wal_fallback_eligible("locking protocol"));
assert!(!wal_fallback_eligible("syntax error"));
}
#[test]
fn open_creates_parent_and_metadata() {
let dir = tempfile::tempdir().unwrap();
let path = SqliteBackend::default_db_path(dir.path());
let backend = SqliteBackend::open(&path).unwrap();
assert!(path.exists());
assert_eq!(backend.list_tables().unwrap().len(), 0);
}
#[test]
fn value_ref_blob_becomes_base64() {
assert_eq!(
value_ref_to_json(ValueRef::Blob(b"foobar")),
serde_json::Value::String("Zm9vYmFy".to_string())
);
assert_eq!(value_ref_to_json(ValueRef::Null), serde_json::Value::Null);
}
#[test]
fn value_ref_real_finite_becomes_number() {
assert_eq!(
value_ref_to_json(ValueRef::Real(1.5)),
serde_json::json!(1.5)
);
}
#[test]
fn value_ref_real_non_finite_becomes_null() {
for f in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
assert_eq!(
value_ref_to_json(ValueRef::Real(f)),
serde_json::Value::Null,
"non-finite real {f} should degrade to JSON null"
);
}
}
}