building out the components, updating the UI in overview and adding a little sqlite viewer
This commit is contained in:
itsamejms
Generated
+136
-2
@@ -752,6 +752,7 @@ name = "dm-pal"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"anyhow",
|
||||
"futures-util",
|
||||
"log",
|
||||
"reqwest 0.12.28",
|
||||
"rusqlite",
|
||||
@@ -759,6 +760,7 @@ dependencies = [
|
||||
"serde_json",
|
||||
"tauri",
|
||||
"tauri-build",
|
||||
"tauri-plugin-dialog",
|
||||
"tauri-plugin-fs",
|
||||
"tauri-plugin-log",
|
||||
"tauri-plugin-store",
|
||||
@@ -2831,6 +2833,7 @@ dependencies = [
|
||||
"bytes",
|
||||
"encoding_rs",
|
||||
"futures-core",
|
||||
"futures-util",
|
||||
"h2",
|
||||
"http",
|
||||
"http-body",
|
||||
@@ -2852,12 +2855,14 @@ dependencies = [
|
||||
"sync_wrapper",
|
||||
"tokio",
|
||||
"tokio-native-tls",
|
||||
"tokio-util",
|
||||
"tower",
|
||||
"tower-http",
|
||||
"tower-service",
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams 0.4.2",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
@@ -2891,10 +2896,34 @@ dependencies = [
|
||||
"url",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"wasm-streams",
|
||||
"wasm-streams 0.5.0",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "rfd"
|
||||
version = "0.16.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a15ad77d9e70a92437d8f74c35d99b4e4691128df018833e99f90bcd36152672"
|
||||
dependencies = [
|
||||
"block2",
|
||||
"dispatch2",
|
||||
"glib-sys",
|
||||
"gobject-sys",
|
||||
"gtk-sys",
|
||||
"js-sys",
|
||||
"log",
|
||||
"objc2",
|
||||
"objc2-app-kit",
|
||||
"objc2-core-foundation",
|
||||
"objc2-foundation",
|
||||
"raw-window-handle",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
"windows-sys 0.60.2",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ring"
|
||||
version = "0.17.14"
|
||||
@@ -3781,6 +3810,24 @@ dependencies = [
|
||||
"walkdir",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-plugin-dialog"
|
||||
version = "2.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "65981abb771e74e571a38196c3baa11c459379164791eba0e67abc1a5fac9884"
|
||||
dependencies = [
|
||||
"log",
|
||||
"raw-window-handle",
|
||||
"rfd",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"tauri",
|
||||
"tauri-plugin",
|
||||
"tauri-plugin-fs",
|
||||
"thiserror 2.0.18",
|
||||
"url",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tauri-plugin-fs"
|
||||
version = "2.5.1"
|
||||
@@ -4607,6 +4654,19 @@ dependencies = [
|
||||
"unicode-ident",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65"
|
||||
dependencies = [
|
||||
"futures-util",
|
||||
"js-sys",
|
||||
"wasm-bindgen",
|
||||
"wasm-bindgen-futures",
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "wasm-streams"
|
||||
version = "0.5.0"
|
||||
@@ -4945,6 +5005,15 @@ dependencies = [
|
||||
"windows-targets 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.60.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb"
|
||||
dependencies = [
|
||||
"windows-targets 0.53.5",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-sys"
|
||||
version = "0.61.2"
|
||||
@@ -4978,13 +5047,30 @@ dependencies = [
|
||||
"windows_aarch64_gnullvm 0.52.6",
|
||||
"windows_aarch64_msvc 0.52.6",
|
||||
"windows_i686_gnu 0.52.6",
|
||||
"windows_i686_gnullvm",
|
||||
"windows_i686_gnullvm 0.52.6",
|
||||
"windows_i686_msvc 0.52.6",
|
||||
"windows_x86_64_gnu 0.52.6",
|
||||
"windows_x86_64_gnullvm 0.52.6",
|
||||
"windows_x86_64_msvc 0.52.6",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-targets"
|
||||
version = "0.53.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3"
|
||||
dependencies = [
|
||||
"windows-link 0.2.1",
|
||||
"windows_aarch64_gnullvm 0.53.1",
|
||||
"windows_aarch64_msvc 0.53.1",
|
||||
"windows_i686_gnu 0.53.1",
|
||||
"windows_i686_gnullvm 0.53.1",
|
||||
"windows_i686_msvc 0.53.1",
|
||||
"windows_x86_64_gnu 0.53.1",
|
||||
"windows_x86_64_gnullvm 0.53.1",
|
||||
"windows_x86_64_msvc 0.53.1",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "windows-threading"
|
||||
version = "0.1.0"
|
||||
@@ -5015,6 +5101,12 @@ version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_gnullvm"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_msvc"
|
||||
version = "0.42.2"
|
||||
@@ -5027,6 +5119,12 @@ version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469"
|
||||
|
||||
[[package]]
|
||||
name = "windows_aarch64_msvc"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnu"
|
||||
version = "0.42.2"
|
||||
@@ -5039,12 +5137,24 @@ version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnu"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnullvm"
|
||||
version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_gnullvm"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_msvc"
|
||||
version = "0.42.2"
|
||||
@@ -5057,6 +5167,12 @@ version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66"
|
||||
|
||||
[[package]]
|
||||
name = "windows_i686_msvc"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnu"
|
||||
version = "0.42.2"
|
||||
@@ -5069,6 +5185,12 @@ version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnu"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnullvm"
|
||||
version = "0.42.2"
|
||||
@@ -5081,6 +5203,12 @@ version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_gnullvm"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_msvc"
|
||||
version = "0.42.2"
|
||||
@@ -5093,6 +5221,12 @@ version = "0.52.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec"
|
||||
|
||||
[[package]]
|
||||
name = "windows_x86_64_msvc"
|
||||
version = "0.53.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650"
|
||||
|
||||
[[package]]
|
||||
name = "winnow"
|
||||
version = "0.5.40"
|
||||
|
||||
@@ -24,6 +24,8 @@ tauri = { version = "2.11.3", features = [] }
|
||||
tauri-plugin-log = "2"
|
||||
tauri-plugin-store = "2"
|
||||
tauri-plugin-fs = "2"
|
||||
reqwest = { version = "0.12", features = ["json"] }
|
||||
tauri-plugin-dialog = "2"
|
||||
reqwest = { version = "0.12", features = ["json", "stream"] }
|
||||
futures-util = "0.3"
|
||||
tokio = { version = "1", features = ["full"] }
|
||||
rusqlite = { version = "0.31", features = ["bundled"] }
|
||||
|
||||
@@ -19,6 +19,8 @@
|
||||
"fs:allow-rename",
|
||||
"fs:allow-read-dir",
|
||||
"fs:allow-appdata-read-recursive",
|
||||
"fs:allow-appdata-write-recursive"
|
||||
"fs:allow-appdata-write-recursive",
|
||||
"dialog:default",
|
||||
"dialog:allow-open"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
use crate::llm::AppState;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::{Path, PathBuf};
|
||||
use tauri::Manager;
|
||||
use tauri_plugin_store::StoreExt;
|
||||
|
||||
// ponytail: the data-dir pref store + key are defined in lib.rs; mirror them
|
||||
// here so this module is self-contained for reads/writes.
|
||||
const PREFS_STORE: &str = "dm-pal-prefs.json";
|
||||
const DATA_DIR_KEY: &str = "dataDir";
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct DataDirInfo {
|
||||
/// Absolute path currently in use (may differ from `configured` if it was
|
||||
/// unset/fell back to the default).
|
||||
pub current: String,
|
||||
/// Absolute path that will be used after the next restart, if different.
|
||||
pub configured: Option<String>,
|
||||
/// The OS-default location ($APPDATA/dm-toolkit).
|
||||
pub default: String,
|
||||
}
|
||||
|
||||
/// Return the current data dir + the configured (pending-restart) value, so
|
||||
/// Settings can show "using X — restart to apply Y".
|
||||
#[tauri::command]
|
||||
pub fn get_data_dir(
|
||||
app: tauri::AppHandle,
|
||||
state: tauri::State<'_, AppState>,
|
||||
) -> Result<DataDirInfo, String> {
|
||||
let default = app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| e.to_string())?
|
||||
.join("dm-toolkit");
|
||||
let configured = app
|
||||
.store(PREFS_STORE)
|
||||
.map_err(|e| e.to_string())?
|
||||
.get(DATA_DIR_KEY)
|
||||
.and_then(|v| v.as_str().map(|s| s.to_string()));
|
||||
Ok(DataDirInfo {
|
||||
current: state.data_dir.display().to_string(),
|
||||
configured,
|
||||
default: default.display().to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Subdirs/files that make up the campaign data and must move together.
|
||||
const DATA_RELATIVE: &[&str] = &["lore/lore.db", "generations.db", "images"];
|
||||
|
||||
/// Recursively copy `src` to `dst`, but never overwrite an existing file —
|
||||
/// a re-locate must not clobber data already at the destination, and never
|
||||
/// deletes the source (the old location stays as a backup until the DM deletes it).
|
||||
/// When both src and dst are directories, recurse so new files inside an
|
||||
/// existing destination tree still get copied.
|
||||
fn copy_if_absent(src: &Path, dst: &Path) -> std::io::Result<()> {
|
||||
if !src.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
if src.is_dir() {
|
||||
std::fs::create_dir_all(dst)?;
|
||||
for entry in std::fs::read_dir(src)? {
|
||||
let entry = entry?;
|
||||
copy_if_absent(&entry.path(), &dst.join(entry.file_name()))?;
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
// leaf file: skip if the destination already has it (no overwrite).
|
||||
if dst.exists() {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(parent) = dst.parent() {
|
||||
std::fs::create_dir_all(parent)?;
|
||||
}
|
||||
std::fs::copy(src, dst)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Set the data dir to `new_dir`: create it, copy existing campaign data
|
||||
/// (lore.db, generations.db, images/) into it, and persist the preference so
|
||||
/// the next restart uses it. The currently-open SQLite connections stay on the
|
||||
/// old paths until restart — the caller must prompt the user to relaunch.
|
||||
#[tauri::command]
|
||||
pub fn set_data_dir(
|
||||
app: tauri::AppHandle,
|
||||
state: tauri::State<'_, AppState>,
|
||||
new_dir: String,
|
||||
) -> Result<DataDirInfo, String> {
|
||||
let new_path = PathBuf::from(new_dir.trim());
|
||||
if !new_path.is_absolute() {
|
||||
return Err("Choose an absolute folder path".into());
|
||||
}
|
||||
std::fs::create_dir_all(&new_path).map_err(|e| format!("can't create dir: {e}"))?;
|
||||
|
||||
// Copy existing data into the new location (no overwrite; old data stays put).
|
||||
for rel in DATA_RELATIVE {
|
||||
let src = state.data_dir.join(rel);
|
||||
let dst = new_path.join(rel);
|
||||
if src.exists() {
|
||||
copy_if_absent(&src, &dst).map_err(|e| format!("migrate {rel}: {e}"))?;
|
||||
}
|
||||
}
|
||||
|
||||
// Persist the preference for the next startup.
|
||||
let store = app.store(PREFS_STORE).map_err(|e| e.to_string())?;
|
||||
store.set(DATA_DIR_KEY, serde_json::Value::String(new_path.display().to_string()));
|
||||
store.save().map_err(|e| e.to_string())?;
|
||||
|
||||
Ok(DataDirInfo {
|
||||
current: state.data_dir.display().to_string(),
|
||||
configured: Some(new_path.display().to_string()),
|
||||
default: app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| e.to_string())?
|
||||
.join("dm-toolkit")
|
||||
.display()
|
||||
.to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Clear the preference so the next restart uses the OS default. Does not
|
||||
/// delete any files — the DM re-points, then restarts.
|
||||
#[tauri::command]
|
||||
pub fn reset_data_dir(app: tauri::AppHandle) -> Result<(), String> {
|
||||
let store = app.store(PREFS_STORE).map_err(|e| e.to_string())?;
|
||||
store.delete(DATA_DIR_KEY);
|
||||
store.save().map_err(|e| e.to_string())?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ─── SQLite viewer (read-only) ─────────────────────────────────
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct SqlQueryReq {
|
||||
/// "lore" or "generations".
|
||||
pub db: String,
|
||||
pub sql: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct SqlQueryResult {
|
||||
pub columns: Vec<String>,
|
||||
pub rows: Vec<Vec<serde_json::Value>>,
|
||||
}
|
||||
|
||||
/// Run a read-only SQL query against one of the two SQLite DBs so a DM can
|
||||
/// inspect the backend from Settings → Advanced. Only SELECT/PRAGMA/WITH.
|
||||
#[tauri::command]
|
||||
pub fn sql_query(state: tauri::State<'_, AppState>, req: SqlQueryReq) -> Result<SqlQueryResult, String> {
|
||||
let res = match req.db.as_str() {
|
||||
"lore" => state.rag.raw_query(&req.sql),
|
||||
"generations" => state.gen.raw_query(&req.sql),
|
||||
other => Err(anyhow::anyhow!("unknown db: {other}")),
|
||||
}
|
||||
.map_err(|e| e.to_string())?;
|
||||
Ok(SqlQueryResult { columns: res.0, rows: res.1 })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn copy_if_absent_does_not_overwrite() {
|
||||
let tmp = std::env::temp_dir().join(format!("dm-pal-data-test-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
let src = tmp.join("src");
|
||||
let dst = tmp.join("dst");
|
||||
std::fs::create_dir_all(src.join("images")).unwrap();
|
||||
std::fs::write(src.join("images/a.png"), b"png").unwrap();
|
||||
std::fs::create_dir_all(src.join("lore")).unwrap();
|
||||
std::fs::write(src.join("lore/lore.db"), b"db").unwrap();
|
||||
|
||||
// first copy creates dst with the files
|
||||
copy_if_absent(&src, &dst).unwrap();
|
||||
assert!(dst.join("images/a.png").exists());
|
||||
assert!(dst.join("lore/lore.db").exists());
|
||||
|
||||
// mutate src; second copy must NOT overwrite existing dst files
|
||||
std::fs::write(src.join("images/a.png"), b"changed").unwrap();
|
||||
copy_if_absent(&src, &dst).unwrap();
|
||||
assert_eq!(std::fs::read(dst.join("images/a.png")).unwrap(), b"png");
|
||||
|
||||
// a new file in src does get copied (absent in dst)
|
||||
std::fs::write(src.join("images/b.png"), b"png2").unwrap();
|
||||
copy_if_absent(&src, &dst).unwrap();
|
||||
assert!(dst.join("images/b.png").exists());
|
||||
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn copy_if_absent_missing_src_is_noop() {
|
||||
let tmp = std::env::temp_dir().join(format!("dm-pal-data-test2-{}", std::process::id()));
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
copy_if_absent(&tmp.join("nope"), &tmp.join("dst")).unwrap();
|
||||
assert!(!tmp.join("dst").exists());
|
||||
let _ = std::fs::remove_dir_all(&tmp);
|
||||
}
|
||||
}
|
||||
@@ -1,8 +1,9 @@
|
||||
use crate::llm::AppState;
|
||||
use serde::Deserialize;
|
||||
use futures_util::StreamExt;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
use std::hash::{Hash, Hasher};
|
||||
use tauri::Manager;
|
||||
use tauri::ipc::Channel;
|
||||
|
||||
// ponytail: DefaultHasher is fine for a cache filename — not crypto, just a stable key.
|
||||
|
||||
@@ -14,14 +15,57 @@ pub struct ImageRequest {
|
||||
}
|
||||
|
||||
/// One line of Ollama's NDJSON image-generation response.
|
||||
#[derive(Debug, Deserialize)]
|
||||
/// `step`/`total` carry progress; the final `done: true` line carries the
|
||||
/// singular `image` base64 field. All fields optional per-line.
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
struct OllamaImageLine {
|
||||
#[serde(default)]
|
||||
done: bool,
|
||||
#[serde(default)]
|
||||
step: Option<u32>,
|
||||
#[serde(default)]
|
||||
total: Option<u32>,
|
||||
#[serde(default)]
|
||||
image: Option<String>,
|
||||
}
|
||||
|
||||
/// Parsed view of one NDJSON line used by both the streaming and the
|
||||
/// buffered paths. Extracted so the parsing logic is unit-testable without
|
||||
/// touching the network.
|
||||
#[derive(Debug, PartialEq)]
|
||||
struct ImageProgress {
|
||||
done: bool,
|
||||
step: Option<u32>,
|
||||
total: Option<u32>,
|
||||
image: Option<String>,
|
||||
}
|
||||
|
||||
fn parse_image_line(line: &str) -> Option<ImageProgress> {
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let parsed: OllamaImageLine = serde_json::from_str(line).ok()?;
|
||||
Some(ImageProgress {
|
||||
done: parsed.done,
|
||||
step: parsed.step,
|
||||
total: parsed.total,
|
||||
image: parsed.image,
|
||||
})
|
||||
}
|
||||
|
||||
/// Channel events for streaming image generation.
|
||||
#[derive(Clone, Serialize)]
|
||||
#[serde(tag = "type", content = "data", rename_all = "camelCase")]
|
||||
pub enum ImageEvent {
|
||||
/// Progress update: (step, total). Either may be None if Ollama omits it.
|
||||
Progress { step: Option<u32>, total: Option<u32> },
|
||||
/// Final result: a `data:image/png;base64,...` URL ready for `<img src>`.
|
||||
Done(String),
|
||||
/// Fatal error.
|
||||
Error(String),
|
||||
}
|
||||
|
||||
/// Generate (or fetch from disk cache) an image for `prompt` via the configured
|
||||
/// Ollama image model. Returns a `data:image/png;base64,...` URL ready for `<img src>`.
|
||||
///
|
||||
@@ -29,7 +73,6 @@ struct OllamaImageLine {
|
||||
/// so the front-end can fall back to a placeholder instead of a confusing timeout.
|
||||
#[tauri::command]
|
||||
pub async fn generate_image(
|
||||
app: tauri::AppHandle,
|
||||
state: tauri::State<'_, AppState>,
|
||||
req: ImageRequest,
|
||||
) -> Result<String, String> {
|
||||
@@ -40,34 +83,92 @@ pub async fn generate_image(
|
||||
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
|
||||
let model = req.model.unwrap_or(config.image_model.clone());
|
||||
|
||||
let cache_dir = app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.map_err(|e| e.to_string())?
|
||||
.join("dm-toolkit")
|
||||
.join("images");
|
||||
std::fs::create_dir_all(&cache_dir).map_err(|e| e.to_string())?;
|
||||
|
||||
// Stable cache key over model + prompt. Regenerating with a tweaked prompt
|
||||
// produces a new file; identical prompt reuses the cached PNG.
|
||||
let mut hasher = DefaultHasher::new();
|
||||
model.hash(&mut hasher);
|
||||
req.prompt.hash(&mut hasher);
|
||||
let cache_path = cache_dir.join(format!("{:016x}.png", hasher.finish()));
|
||||
|
||||
// Cache hit: return the stored PNG without calling the model.
|
||||
if cache_path.exists() {
|
||||
let (cache_path, cache_hit) = prepare_cache(&state.data_dir, &model, &req.prompt)?;
|
||||
if cache_hit {
|
||||
let bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
|
||||
return Ok(data_url(&bytes));
|
||||
}
|
||||
|
||||
let png_bytes = request_image_bytes(&config, &model, &req.prompt, None).await?;
|
||||
std::fs::write(&cache_path, &png_bytes).map_err(|e| e.to_string())?;
|
||||
Ok(data_url(&png_bytes))
|
||||
}
|
||||
|
||||
/// Streaming variant: emits `ImageEvent::Progress` as Ollama reports `step`/`total`,
|
||||
/// then `ImageEvent::Done` with the data URL (or `Error`). Reuses the same disk cache
|
||||
/// as `generate_image`. The DM gets a real progress bar for the multi-second wait.
|
||||
#[tauri::command]
|
||||
pub async fn generate_image_stream(
|
||||
state: tauri::State<'_, AppState>,
|
||||
req: ImageRequest,
|
||||
channel: Channel<ImageEvent>,
|
||||
) -> Result<(), String> {
|
||||
if cfg!(not(target_os = "macos")) {
|
||||
let _ = channel.send(ImageEvent::Error(
|
||||
"image generation is macOS-only via Ollama (for now)".into(),
|
||||
));
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let config = state.config.lock().map_err(|e| e.to_string())?.clone();
|
||||
let model = req.model.unwrap_or(config.image_model.clone());
|
||||
|
||||
let (cache_path, cache_hit) = prepare_cache(&state.data_dir, &model, &req.prompt)?;
|
||||
if cache_hit {
|
||||
let bytes = std::fs::read(&cache_path).map_err(|e| e.to_string())?;
|
||||
let _ = channel.send(ImageEvent::Done(data_url(&bytes)));
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Drive the request on a background task so the command returns immediately
|
||||
// and progress flows through the channel. Errors become ImageEvent::Error.
|
||||
let channel = std::sync::Arc::new(channel);
|
||||
let ch = channel.clone();
|
||||
tauri::async_runtime::spawn(async move {
|
||||
match request_image_bytes(&config, &model, &req.prompt, Some(ch)).await {
|
||||
Ok(png_bytes) => {
|
||||
let _ = std::fs::write(&cache_path, &png_bytes);
|
||||
let _ = channel.send(ImageEvent::Done(data_url(&png_bytes)));
|
||||
}
|
||||
Err(e) => {
|
||||
let _ = channel.send(ImageEvent::Error(e));
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Resolve the on-disk cache path for (model, prompt) and report a cache hit.
|
||||
/// `base` is the configured data dir (AppState.data_dir).
|
||||
fn prepare_cache(
|
||||
base: &std::path::Path,
|
||||
model: &str,
|
||||
prompt: &str,
|
||||
) -> Result<(std::path::PathBuf, bool), String> {
|
||||
let cache_dir = base.join("images");
|
||||
std::fs::create_dir_all(&cache_dir).map_err(|e| e.to_string())?;
|
||||
|
||||
let mut hasher = DefaultHasher::new();
|
||||
model.hash(&mut hasher);
|
||||
prompt.hash(&mut hasher);
|
||||
let cache_path = cache_dir.join(format!("{:016x}.png", hasher.finish()));
|
||||
let hit = cache_path.exists();
|
||||
Ok((cache_path, hit))
|
||||
}
|
||||
|
||||
/// POST the image request to Ollama and collect the PNG bytes. When a channel
|
||||
/// is given, parse the NDJSON body incrementally and emit progress; otherwise
|
||||
/// read the whole body at once (legacy buffered path).
|
||||
async fn request_image_bytes(
|
||||
config: &crate::llm::LlmConfig,
|
||||
model: &str,
|
||||
prompt: &str,
|
||||
channel: Option<std::sync::Arc<Channel<ImageEvent>>>,
|
||||
) -> Result<Vec<u8>, String> {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/api/generate", config.api_url.trim_end_matches('/'));
|
||||
let body = serde_json::json!({
|
||||
"model": model,
|
||||
"prompt": req.prompt,
|
||||
"stream": false,
|
||||
});
|
||||
let body = serde_json::json!({ "model": model, "prompt": prompt, "stream": false });
|
||||
|
||||
let res = client
|
||||
.post(&url)
|
||||
@@ -82,19 +183,47 @@ pub async fn generate_image(
|
||||
return Err(format!("image error {status}: {text}"));
|
||||
}
|
||||
|
||||
// Ollama returns newline-delimited JSON even with stream:false for image models;
|
||||
// the final line with `done: true` carries the singular `image` base64 field.
|
||||
let body_text = res.text().await.map_err(|e| format!("image read error: {e}"))?;
|
||||
// ponytail: Ollama image models emit NDJSON even with stream:false —
|
||||
// progress lines (step/total) then a final done:true carrying the image.
|
||||
// The buffered path reads it all at once; the streaming path splits lines
|
||||
// as they arrive so the bar animates.
|
||||
let mut png_b64: Option<String> = None;
|
||||
for line in body_text.lines() {
|
||||
let line = line.trim();
|
||||
if line.is_empty() {
|
||||
continue;
|
||||
let mut buffer = String::new();
|
||||
|
||||
let apply_line = |line: &str, png: &mut Option<String>, ch: Option<&Channel<ImageEvent>>| {
|
||||
if let Some(p) = parse_image_line(line) {
|
||||
if let Some(c) = ch {
|
||||
let _ = c.send(ImageEvent::Progress { step: p.step, total: p.total });
|
||||
}
|
||||
if let Some(b64) = p.image {
|
||||
*png = Some(b64);
|
||||
}
|
||||
}
|
||||
if let Ok(parsed) = serde_json::from_str::<OllamaImageLine>(line) {
|
||||
if let Some(b64) = parsed.image {
|
||||
png_b64 = Some(b64);
|
||||
if parsed.done {
|
||||
};
|
||||
|
||||
if let Some(ch) = channel {
|
||||
let mut stream = res.bytes_stream();
|
||||
while let Some(chunk) = stream.next().await {
|
||||
let chunk = chunk.map_err(|e| format!("image read error: {e}"))?;
|
||||
buffer.push_str(&String::from_utf8_lossy(&chunk));
|
||||
// Process complete lines; keep the trailing partial line in buffer.
|
||||
while let Some(idx) = buffer.find('\n') {
|
||||
let line = buffer.split_off(idx + 1);
|
||||
let complete = std::mem::replace(&mut buffer, line);
|
||||
apply_line(&complete, &mut png_b64, Some(&ch));
|
||||
}
|
||||
}
|
||||
if !buffer.trim().is_empty() {
|
||||
apply_line(&buffer, &mut png_b64, Some(&ch));
|
||||
}
|
||||
} else {
|
||||
let body_text = res.text().await.map_err(|e| format!("image read error: {e}"))?;
|
||||
for line in body_text.lines() {
|
||||
apply_line(line, &mut png_b64, None);
|
||||
// ponytail: stop after the done line in the buffered path — the
|
||||
// final image is the one we want.
|
||||
if let Some(p) = parse_image_line(line) {
|
||||
if p.done && p.image.is_some() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
@@ -102,11 +231,8 @@ pub async fn generate_image(
|
||||
}
|
||||
|
||||
let b64 = png_b64.ok_or_else(|| "no image data in Ollama response".to_string())?;
|
||||
|
||||
// Decode + persist to cache, then return a data URL.
|
||||
let png_bytes = base64_decode(&b64)?;
|
||||
std::fs::write(&cache_path, &png_bytes).map_err(|e| e.to_string())?;
|
||||
Ok(data_url(&png_bytes))
|
||||
Ok(png_bytes)
|
||||
}
|
||||
|
||||
fn data_url(png: &[u8]) -> String {
|
||||
@@ -115,7 +241,7 @@ fn data_url(png: &[u8]) -> String {
|
||||
let mut out = String::with_capacity((png.len() + 2) / 3 * 4);
|
||||
let mut chunks = png.chunks_exact(3);
|
||||
for c in &mut chunks {
|
||||
let n = (c[0] as usize) << 16 | (c[1] as usize) << 8 | c[2] as usize;
|
||||
let n = (c[0] as usize) << 16 | (c[1] as usize) << 8 | (c[2] as usize);
|
||||
out.push(table[(n >> 18) & 63] as char);
|
||||
out.push(table[(n >> 12) & 63] as char);
|
||||
out.push(table[(n >> 6) & 63] as char);
|
||||
@@ -200,4 +326,48 @@ mod tests {
|
||||
let s = format!("{:016x}", h.finish());
|
||||
assert_eq!(s.len(), 16);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_progress_and_image_lines() {
|
||||
// intermediate progress line, no image
|
||||
let p = parse_image_line(r#"{"model":"x/flux2-klein:4b","done":false,"total":4,"step":1}"#).unwrap();
|
||||
assert_eq!(p, ImageProgress { done: false, step: Some(1), total: Some(4), image: None });
|
||||
|
||||
// final done line carries the image
|
||||
let p = parse_image_line(r#"{"model":"x/flux2-klein:4b","done":true,"image":"iVBORw0KGgoAAAANSUhEUgAA"}"#).unwrap();
|
||||
assert!(p.done);
|
||||
assert_eq!(p.image.as_deref(), Some("iVBORw0KGgoAAAANSUhEUgAA"));
|
||||
assert!(p.total.is_none() && p.step.is_none());
|
||||
|
||||
// blank/garbage lines are ignored, not errors
|
||||
assert!(parse_image_line("").is_none());
|
||||
assert!(parse_image_line("not json").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn split_ndjson_buffer_keeps_trailing_partial() {
|
||||
// Simulate two chunks arriving separately where the split falls mid-line.
|
||||
let chunk1 = "{\"done\":false,\"step\":1,\"total\":4}\n{\"done\":tru";
|
||||
let chunk2 = "e,\"image\":\"abc\"}\n";
|
||||
let mut buffer = String::new();
|
||||
let mut png: Option<String> = None;
|
||||
let whole = format!("{chunk1}{chunk2}");
|
||||
// emulate the streaming loop over the concatenated body
|
||||
buffer.push_str(&whole);
|
||||
let mut lines = Vec::new();
|
||||
while let Some(idx) = buffer.find('\n') {
|
||||
let line = buffer.split_off(idx + 1);
|
||||
let complete = std::mem::replace(&mut buffer, line);
|
||||
lines.push(complete);
|
||||
}
|
||||
if !buffer.trim().is_empty() {
|
||||
lines.push(std::mem::take(&mut buffer));
|
||||
}
|
||||
for line in &lines {
|
||||
if let Some(p) = parse_image_line(line) {
|
||||
if let Some(b) = p.image { png = Some(b); }
|
||||
}
|
||||
}
|
||||
assert_eq!(png.as_deref(), Some("abc"));
|
||||
}
|
||||
}
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod llm_commands;
|
||||
pub mod image_commands;
|
||||
pub mod rag_commands;
|
||||
pub mod generation_commands;
|
||||
pub mod generation_commands;
|
||||
pub mod data_commands;
|
||||
@@ -55,4 +55,23 @@ pub fn rag_list(state: tauri::State<'_, AppState>) -> Result<Vec<RagSource>, Str
|
||||
#[tauri::command]
|
||||
pub fn rag_clear(state: tauri::State<'_, AppState>, source: Option<String>) -> Result<(), String> {
|
||||
state.rag.clear(source.as_deref()).map_err(|e| e.to_string())
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct RagChunk {
|
||||
pub id: i64,
|
||||
pub preview: String,
|
||||
}
|
||||
|
||||
/// Preview the chunks stored under `source` (first 200 chars each). Lets the
|
||||
/// DM verify what actually got embedded without dumping the whole corpus.
|
||||
#[tauri::command]
|
||||
pub fn rag_chunks(state: tauri::State<'_, AppState>, source: String) -> Result<Vec<RagChunk>, String> {
|
||||
state
|
||||
.rag
|
||||
.list_chunks(&source)
|
||||
.map_err(|e| e.to_string())?
|
||||
.into_iter()
|
||||
.map(|(id, preview)| Ok(RagChunk { id, preview }))
|
||||
.collect()
|
||||
}
|
||||
@@ -144,6 +144,15 @@ impl GenerationStore {
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Read-only SQL pass-through for the advanced SQLite viewer.
|
||||
pub fn raw_query(
|
||||
&self,
|
||||
sql: &str,
|
||||
) -> anyhow::Result<(Vec<String>, Vec<Vec<serde_json::Value>>)> {
|
||||
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
|
||||
crate::sql_viewer::run_readonly_query(&db, sql)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
+36
-5
@@ -2,10 +2,37 @@ mod llm;
|
||||
mod rag;
|
||||
mod generations;
|
||||
mod commands;
|
||||
mod sql_viewer;
|
||||
|
||||
use llm::AppState;
|
||||
use std::sync::Mutex;
|
||||
use tauri::Manager;
|
||||
use tauri_plugin_store::StoreExt;
|
||||
|
||||
// ponytail: the data-dir preference lives in a tiny store at the OS default
|
||||
// app_data_dir so it's always discoverable at startup, even before we know the
|
||||
// custom location. Key: `dataDir` (absolute path). Empty/missing → default.
|
||||
const PREFS_STORE: &str = "dm-pal-prefs.json";
|
||||
const DATA_DIR_KEY: &str = "dataDir";
|
||||
|
||||
/// Resolve the data dir: the configured one if set and usable, else the
|
||||
/// default `$APPDATA/dm-toolkit`.
|
||||
fn resolve_data_dir(app: &tauri::AppHandle) -> std::path::PathBuf {
|
||||
let default = app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.expect("app data dir")
|
||||
.join("dm-toolkit");
|
||||
if let Ok(store) = app.store(PREFS_STORE) {
|
||||
if let Some(Some(v)) = store.get(DATA_DIR_KEY).map(|v| v.as_str().map(|s| s.to_string())) {
|
||||
let p = std::path::PathBuf::from(v);
|
||||
// ponytail: fall back to default if the configured path is gone —
|
||||
// a DM who unmounts the external drive shouldn't lose the app.
|
||||
if p.is_absolute() { return p; }
|
||||
}
|
||||
}
|
||||
default
|
||||
}
|
||||
|
||||
#[cfg_attr(mobile, tauri::mobile_entry_point)]
|
||||
pub fn run() {
|
||||
@@ -13,12 +40,9 @@ pub fn run() {
|
||||
.plugin(tauri_plugin_log::Builder::default().build())
|
||||
.plugin(tauri_plugin_store::Builder::default().build())
|
||||
.plugin(tauri_plugin_fs::init())
|
||||
.plugin(tauri_plugin_dialog::init())
|
||||
.setup(|app| {
|
||||
let data_dir = app
|
||||
.path()
|
||||
.app_data_dir()
|
||||
.expect("app data dir")
|
||||
.join("dm-toolkit");
|
||||
let data_dir = resolve_data_dir(&app.handle());
|
||||
let rag = rag::RagStore::open(&data_dir.join("lore"))
|
||||
.expect("open lore db");
|
||||
let gen = generations::GenerationStore::open(&data_dir)
|
||||
@@ -27,6 +51,7 @@ pub fn run() {
|
||||
config: Mutex::new(llm::LlmConfig::default()),
|
||||
rag,
|
||||
gen,
|
||||
data_dir,
|
||||
});
|
||||
Ok(())
|
||||
})
|
||||
@@ -38,15 +63,21 @@ pub fn run() {
|
||||
commands::llm_commands::set_llm_config,
|
||||
commands::llm_commands::test_connection,
|
||||
commands::image_commands::generate_image,
|
||||
commands::image_commands::generate_image_stream,
|
||||
commands::rag_commands::rag_add,
|
||||
commands::rag_commands::rag_search,
|
||||
commands::rag_commands::rag_list,
|
||||
commands::rag_commands::rag_clear,
|
||||
commands::rag_commands::rag_chunks,
|
||||
commands::generation_commands::generation_add,
|
||||
commands::generation_commands::generation_list,
|
||||
commands::generation_commands::generation_get,
|
||||
commands::generation_commands::generation_delete,
|
||||
commands::generation_commands::generation_counts,
|
||||
commands::data_commands::get_data_dir,
|
||||
commands::data_commands::set_data_dir,
|
||||
commands::data_commands::reset_data_dir,
|
||||
commands::data_commands::sql_query,
|
||||
])
|
||||
.run(tauri::generate_context!())
|
||||
.expect("error while running tauri application");
|
||||
|
||||
@@ -17,6 +17,10 @@ pub struct AppState {
|
||||
pub config: Mutex<LlmConfig>,
|
||||
pub rag: crate::rag::RagStore,
|
||||
pub gen: crate::generations::GenerationStore,
|
||||
/// Where lore.db, generations.db, and cached images live. Configurable via
|
||||
/// the `dataDir` key in the `dm-pal-prefs.json` store; defaults to
|
||||
/// `$APPDATA/dm-toolkit` when unset. Set once at startup.
|
||||
pub data_dir: std::path::PathBuf,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
|
||||
@@ -179,6 +179,30 @@ impl RagStore {
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Preview chunks for a source: (id, first 200 chars of text). Used by the
|
||||
/// Lore panel so a DM can see what actually got embedded without a full dump.
|
||||
pub fn list_chunks(&self, source: &str) -> anyhow::Result<Vec<(i64, String)>> {
|
||||
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
|
||||
let mut stmt = db.prepare("SELECT id, substr(text, 1, 200) FROM chunks WHERE source = ? ORDER BY id")?;
|
||||
let rows = stmt.query_map(params![source], |r| {
|
||||
Ok((r.get::<_, i64>(0)?, r.get::<_, String>(1)?))
|
||||
})?;
|
||||
let mut out = Vec::new();
|
||||
for r in rows {
|
||||
out.push(r?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Read-only SQL pass-through for the advanced SQLite viewer.
|
||||
pub fn raw_query(
|
||||
&self,
|
||||
sql: &str,
|
||||
) -> anyhow::Result<(Vec<String>, Vec<Vec<serde_json::Value>>)> {
|
||||
let db = self.db.lock().map_err(|e| anyhow::anyhow!("db lock: {e}"))?;
|
||||
crate::sql_viewer::run_readonly_query(&db, sql)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
use rusqlite::Connection;
|
||||
use rusqlite::types::Value as SqlValue;
|
||||
use serde_json::Value as JsonValue;
|
||||
|
||||
// ponytail: a read-only SQL viewer for the two SQLite DBs (lore.db,
|
||||
// generations.db). We gate on the statement starting with SELECT / PRAGMA /
|
||||
// WITH so a curious DM can inspect the backend but can't mutate it from the
|
||||
// UI. Enough for "what's going on in the backend" without a footgun.
|
||||
|
||||
/// Run a read-only query against `conn`. Returns (column names, rows).
|
||||
pub fn run_readonly_query(
|
||||
conn: &Connection,
|
||||
sql: &str,
|
||||
) -> anyhow::Result<(Vec<String>, Vec<Vec<JsonValue>>)> {
|
||||
let trimmed = sql.trim().trim_end_matches(';').trim();
|
||||
if trimmed.is_empty() {
|
||||
return Ok((Vec::new(), Vec::new()));
|
||||
}
|
||||
let upper = trimmed.to_uppercase();
|
||||
let allowed = upper.starts_with("SELECT")
|
||||
|| upper.starts_with("PRAGMA")
|
||||
|| upper.starts_with("WITH");
|
||||
if !allowed {
|
||||
anyhow::bail!("Only read-only queries (SELECT / PRAGMA / WITH) are allowed");
|
||||
}
|
||||
|
||||
let mut stmt = conn.prepare(trimmed)?;
|
||||
let col_count = stmt.column_count();
|
||||
let columns: Vec<String> = (0..col_count)
|
||||
.map(|i| stmt.column_name(i).unwrap_or("?").to_string())
|
||||
.collect();
|
||||
let rows = stmt.query_map([], |row| {
|
||||
let mut vals = Vec::with_capacity(col_count);
|
||||
for i in 0..col_count {
|
||||
let v: SqlValue = row.get(i)?;
|
||||
vals.push(sql_to_json(v));
|
||||
}
|
||||
Ok(vals)
|
||||
})?;
|
||||
let mut out = Vec::new();
|
||||
for r in rows {
|
||||
out.push(r?);
|
||||
}
|
||||
Ok((columns, out))
|
||||
}
|
||||
|
||||
fn sql_to_json(v: SqlValue) -> JsonValue {
|
||||
match v {
|
||||
SqlValue::Null => JsonValue::Null,
|
||||
SqlValue::Integer(i) => i.into(),
|
||||
SqlValue::Real(f) => f.into(),
|
||||
SqlValue::Text(s) => JsonValue::String(s),
|
||||
// ponytail: blobs aren't useful to render raw — show a size hint so the
|
||||
// DM knows there's data without dumping megabytes of embedding bytes.
|
||||
SqlValue::Blob(b) => JsonValue::String(format!("<blob {} bytes>", b.len())),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn mem() -> Connection {
|
||||
let c = Connection::open_in_memory().unwrap();
|
||||
c.execute_batch("CREATE TABLE t (id INTEGER PRIMARY KEY, name TEXT, weight REAL);").unwrap();
|
||||
c.execute("INSERT INTO t (name, weight) VALUES (?1, ?2), (?1, ?2)", rusqlite::params!["goblin", 23.5]).unwrap();
|
||||
c
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn select_returns_columns_and_rows() {
|
||||
let (cols, rows) = run_readonly_query(&mem(), "SELECT id, name, weight FROM t ORDER BY id").unwrap();
|
||||
assert_eq!(cols, vec!["id".to_string(), "name".into(), "weight".into()]);
|
||||
assert_eq!(rows.len(), 2);
|
||||
assert_eq!(rows[0][1], serde_json::json!("goblin"));
|
||||
assert_eq!(rows[0][2], serde_json::json!(23.5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_writes() {
|
||||
let c = mem();
|
||||
let err = run_readonly_query(&c, "DELETE FROM t").unwrap_err();
|
||||
assert!(err.to_string().contains("read-only"));
|
||||
let err = run_readonly_query(&c, "INSERT INTO t (name, weight) VALUES ('x', 1)").unwrap_err();
|
||||
assert!(err.to_string().contains("read-only"));
|
||||
let err = run_readonly_query(&c, "UPDATE t SET name='x'").unwrap_err();
|
||||
assert!(err.to_string().contains("read-only"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pragma_table_info_allowed() {
|
||||
let (cols, rows) = run_readonly_query(&mem(), "PRAGMA table_info(t)").unwrap();
|
||||
assert!(cols.contains(&"name".to_string()));
|
||||
assert!(rows.iter().any(|r| r.contains(&serde_json::json!("id"))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn with_cte_allowed() {
|
||||
let (_cols, rows) = run_readonly_query(&mem(), "WITH x AS (SELECT 1 AS n) SELECT * FROM x").unwrap();
|
||||
assert_eq!(rows.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_and_trailing_semicolon_ok() {
|
||||
let (cols, rows) = run_readonly_query(&mem(), "SELECT 1 AS n;").unwrap();
|
||||
assert_eq!(cols, vec!["n"]);
|
||||
assert_eq!(rows.len(), 1);
|
||||
let Ok((c, r)) = run_readonly_query(&mem(), " ") else { panic!("empty query failed") };
|
||||
assert!(c.is_empty() && r.is_empty());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user