perspt_agent/tools/families/
db.rs

1//! Local DB explorer family: read-only, bounded queries over data files the
2//! workspace already contains (CSV, Parquet, JSON, DuckDB).
3//!
4//! `db_query` runs SELECT-only SQL against a fresh **in-memory** DuckDB
5//! connection in which the target file is pre-registered as the view
6//! `data`. The SQL text passes a statement-class allowlist (a single
7//! SELECT/WITH statement) and a blocklist of filesystem- and
8//! environment-reaching functions, so a query can never reach beyond the
9//! validated workspace file. Results carry row and byte caps. The session
10//! store is never a valid target, so DuckDB's one-live-handle constraint
11//! does not bite.
12
13use std::path::Path;
14use std::sync::Arc;
15
16use anyhow::{Context, Result};
17use perspt_sdk::{AccessMode, EffectKind, FootprintSpec, ResourceSelector, ToolEntry};
18
19use super::{family_entry, object_schema};
20use crate::candidate::CandidateWorkspace;
21use crate::toolloop::EffectOutcome;
22use crate::tools::handlers::{CandidateHandlerRegistry, CandidateToolHandler};
23
24const MAX_ROWS: usize = 200;
25const MAX_BYTES: usize = 32 * 1024;
26const DATA_EXTENSIONS: &[&str] = &["csv", "tsv", "parquet", "json", "jsonl", "ndjson", "duckdb"];
27
28fn scoped_path_footprint() -> FootprintSpec {
29    FootprintSpec::new(vec![ResourceSelector::ScopedArgument {
30        family: "db".into(),
31        field: "path".into(),
32        access: AccessMode::Read,
33    }])
34}
35
36pub fn entries() -> Vec<ToolEntry> {
37    vec![
38        family_entry(
39            "db_list",
40            "List data files in the workspace (CSV, TSV, Parquet, JSON, DuckDB)",
41            EffectKind::DataRead,
42            object_schema(&[]),
43            FootprintSpec::new(vec![ResourceSelector::Literal {
44                resource: perspt_sdk::Resource::Scoped {
45                    family: "db".into(),
46                    key: "list".into(),
47                },
48                access: AccessMode::Read,
49            }]),
50        ),
51        family_entry(
52            "db_schema",
53            "Describe the columns and types of a workspace data file",
54            EffectKind::DataRead,
55            object_schema(&[("path", "string", "Workspace-relative data file", true)]),
56            scoped_path_footprint(),
57        ),
58        family_entry(
59            "db_query",
60            "Run one SELECT-only SQL statement against a workspace data file; \
61             the file is exposed as the view `data`; results are row- and \
62             byte-capped",
63            EffectKind::DataRead,
64            object_schema(&[
65                ("path", "string", "Workspace-relative data file", true),
66                (
67                    "sql",
68                    "string",
69                    "A single SELECT/WITH statement over `data`",
70                    true,
71                ),
72            ]),
73            scoped_path_footprint(),
74        ),
75    ]
76}
77
78pub fn register(registry: &mut CandidateHandlerRegistry) -> Result<()> {
79    registry.register("db_list", Arc::new(DbList))?;
80    registry.register("db_schema", Arc::new(DbSchema))?;
81    registry.register("db_query", Arc::new(DbQuery))?;
82    Ok(())
83}
84
85fn is_data_file(path: &Path) -> bool {
86    path.extension()
87        .and_then(|ext| ext.to_str())
88        .map(|ext| DATA_EXTENSIONS.contains(&ext.to_ascii_lowercase().as_str()))
89        .unwrap_or(false)
90}
91
92/// Resolve a validated workspace-relative data-file argument.
93fn data_file(
94    workspace: &CandidateWorkspace,
95    call: &perspt_sdk::ProviderToolCall,
96) -> Result<std::path::PathBuf> {
97    let relative = call
98        .arguments
99        .get("path")
100        .and_then(|v| v.as_str())
101        .context("a workspace-relative `path` argument is required")?;
102    let relative = workspace.validate_relative(relative)?;
103    let path = workspace.overlay_root().join(&relative);
104    anyhow::ensure!(
105        is_data_file(&path),
106        "not a recognized data file (csv, tsv, parquet, json, jsonl, ndjson, duckdb): {relative}"
107    );
108    anyhow::ensure!(path.is_file(), "data file not found: {relative}");
109    Ok(path)
110}
111
112/// Open an in-memory connection with the target file registered as the
113/// read-only view `data`, then lock external access down.
114fn open_data_view(path: &Path) -> Result<duckdb::Connection> {
115    let connection = duckdb::Connection::open_in_memory()?;
116    let literal = path
117        .to_str()
118        .context("data file path is not valid UTF-8")?
119        .replace('\'', "''");
120    let extension = path
121        .extension()
122        .and_then(|ext| ext.to_str())
123        .unwrap_or_default()
124        .to_ascii_lowercase();
125    let view = match extension.as_str() {
126        "csv" | "tsv" => format!("CREATE VIEW data AS SELECT * FROM read_csv_auto('{literal}')"),
127        "parquet" => format!("CREATE VIEW data AS SELECT * FROM read_parquet('{literal}')"),
128        "json" | "jsonl" | "ndjson" => {
129            format!("CREATE VIEW data AS SELECT * FROM read_json_auto('{literal}')")
130        }
131        "duckdb" => format!("ATTACH '{literal}' AS src (READ_ONLY)"),
132        other => anyhow::bail!("unsupported data file extension {other:?}"),
133    };
134    connection.execute_batch(&view)?;
135    Ok(connection)
136}
137
138/// A single SELECT/WITH statement with no filesystem- or environment-
139/// reaching functions. Fail closed on anything else.
140fn validate_select_only(sql: &str) -> Result<()> {
141    let trimmed = sql.trim().trim_end_matches(';');
142    anyhow::ensure!(
143        !trimmed.contains(';'),
144        "db_query admits exactly one statement"
145    );
146    let lowered = trimmed.to_ascii_lowercase();
147    anyhow::ensure!(
148        lowered.starts_with("select") || lowered.starts_with("with"),
149        "db_query admits only SELECT/WITH statements"
150    );
151    const FORBIDDEN: &[&str] = &[
152        "attach",
153        "copy",
154        "install",
155        "load",
156        "pragma",
157        "export",
158        "import",
159        "create",
160        "insert",
161        "update ",
162        "delete",
163        "drop",
164        "alter",
165        "set ",
166        "call",
167        "read_csv",
168        "read_parquet",
169        "read_json",
170        "read_text",
171        "read_blob",
172        "glob",
173        "getenv",
174        "sniff_csv",
175    ];
176    for keyword in FORBIDDEN {
177        anyhow::ensure!(
178            !lowered.contains(keyword),
179            "db_query rejects {} in SQL (read-only, pre-registered view only)",
180            keyword.trim()
181        );
182    }
183    Ok(())
184}
185
186/// Render capped query results as aligned text.
187fn render_rows(statement: &mut duckdb::Statement<'_>) -> Result<String> {
188    let mut rows = statement.query([])?;
189    let mut lines = Vec::new();
190    let mut names_written = false;
191    let mut row_count = 0usize;
192    let mut bytes = 0usize;
193    while let Some(row) = rows.next()? {
194        let stmt = row.as_ref();
195        if !names_written {
196            lines.push(stmt.column_names().join(" | "));
197            names_written = true;
198        }
199        let width = stmt.column_count();
200        let mut cells = Vec::with_capacity(width);
201        for index in 0..width {
202            let value: duckdb::types::Value = row.get(index)?;
203            cells.push(render_value(&value));
204        }
205        let line = cells.join(" | ");
206        bytes += line.len();
207        lines.push(line);
208        row_count += 1;
209        if row_count >= MAX_ROWS || bytes >= MAX_BYTES {
210            lines.push(format!("… (capped at {row_count} rows)"));
211            break;
212        }
213    }
214    if lines.is_empty() {
215        lines.push("(no rows)".into());
216    }
217    Ok(lines.join("\n"))
218}
219
220fn render_value(value: &duckdb::types::Value) -> String {
221    match value {
222        duckdb::types::Value::Null => "NULL".into(),
223        duckdb::types::Value::Text(text) => text.clone(),
224        other => format!("{other:?}"),
225    }
226}
227
228struct DbList;
229
230#[async_trait::async_trait]
231impl CandidateToolHandler for DbList {
232    async fn apply(
233        &self,
234        workspace: &CandidateWorkspace,
235        _call: &perspt_sdk::ProviderToolCall,
236        _entry: &ToolEntry,
237    ) -> Result<EffectOutcome> {
238        let root = workspace.overlay_root().to_path_buf();
239        let mut files = Vec::new();
240        for entry in ignore::WalkBuilder::new(&root).build().flatten() {
241            let path = entry.path();
242            if path.is_file() && is_data_file(path) {
243                if let Ok(relative) = path.strip_prefix(&root) {
244                    let size = entry.metadata().map(|m| m.len()).unwrap_or(0);
245                    files.push(format!("{} ({size} bytes)", relative.display()));
246                }
247            }
248        }
249        files.sort();
250        Ok(EffectOutcome {
251            output: if files.is_empty() {
252                "no data files found".into()
253            } else {
254                files.join("\n")
255            },
256            mutated: false,
257            completed: true,
258        })
259    }
260}
261
262struct DbSchema;
263
264#[async_trait::async_trait]
265impl CandidateToolHandler for DbSchema {
266    async fn apply(
267        &self,
268        workspace: &CandidateWorkspace,
269        call: &perspt_sdk::ProviderToolCall,
270        _entry: &ToolEntry,
271    ) -> Result<EffectOutcome> {
272        let path = data_file(workspace, call)?;
273        let connection = open_data_view(&path)?;
274        let sql = if path.extension().and_then(|e| e.to_str()) == Some("duckdb") {
275            "SELECT table_name, column_name, data_type FROM information_schema.columns \
276             ORDER BY table_name, ordinal_position"
277        } else {
278            "SELECT 'data' AS table_name, column_name, data_type \
279             FROM information_schema.columns WHERE table_name = 'data' \
280             ORDER BY ordinal_position"
281        };
282        let mut statement = connection.prepare(sql)?;
283        let output = render_rows(&mut statement)?;
284        Ok(EffectOutcome {
285            output,
286            mutated: false,
287            completed: true,
288        })
289    }
290}
291
292struct DbQuery;
293
294#[async_trait::async_trait]
295impl CandidateToolHandler for DbQuery {
296    async fn apply(
297        &self,
298        workspace: &CandidateWorkspace,
299        call: &perspt_sdk::ProviderToolCall,
300        _entry: &ToolEntry,
301    ) -> Result<EffectOutcome> {
302        let sql = call
303            .arguments
304            .get("sql")
305            .and_then(|v| v.as_str())
306            .context("db_query requires a `sql` argument")?;
307        validate_select_only(sql)?;
308        let path = data_file(workspace, call)?;
309        let connection = open_data_view(&path)?;
310        let mut statement = connection.prepare(sql.trim().trim_end_matches(';'))?;
311        let output = render_rows(&mut statement)?;
312        Ok(EffectOutcome {
313            output,
314            mutated: false,
315            completed: true,
316        })
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    fn probe_call(name: &str, arguments: serde_json::Value) -> perspt_sdk::ProviderToolCall {
325        perspt_sdk::ProviderToolCall {
326            call_id: "c1".into(),
327            name: name.into(),
328            arguments,
329        }
330    }
331
332    #[test]
333    fn select_only_allowlist_rejects_escapes() {
334        assert!(validate_select_only("SELECT * FROM data").is_ok());
335        assert!(validate_select_only("WITH t AS (SELECT 1) SELECT * FROM t").is_ok());
336        assert!(validate_select_only("SELECT * FROM read_csv_auto('/etc/passwd')").is_err());
337        assert!(validate_select_only("INSTALL httpfs").is_err());
338        assert!(validate_select_only("SELECT 1; DROP TABLE data").is_err());
339        assert!(validate_select_only("PRAGMA database_list").is_err());
340        assert!(validate_select_only("SELECT getenv('HOME')").is_err());
341        assert!(validate_select_only("COPY data TO '/tmp/out.csv'").is_err());
342    }
343
344    #[tokio::test]
345    async fn db_query_reads_workspace_csv_only() {
346        let dir = tempfile::tempdir().unwrap();
347        std::fs::write(
348            dir.path().join("people.csv"),
349            "name,age\nada,36\ngrace,45\n",
350        )
351        .unwrap();
352        let workspace = CandidateWorkspace::create(dir.path(), "n1", 0, "r1").unwrap();
353
354        let outcome = DbQuery
355            .apply(
356                &workspace,
357                &probe_call(
358                    "db_query",
359                    serde_json::json!({
360                        "path": "people.csv",
361                        "sql": "SELECT name FROM data ORDER BY age DESC"
362                    }),
363                ),
364                &entries()[2],
365            )
366            .await
367            .unwrap();
368        assert!(outcome.output.contains("grace"));
369        assert!(!outcome.mutated);
370
371        let escape = DbQuery
372            .apply(
373                &workspace,
374                &probe_call(
375                    "db_query",
376                    serde_json::json!({
377                        "path": "../outside.csv",
378                        "sql": "SELECT * FROM data"
379                    }),
380                ),
381                &entries()[2],
382            )
383            .await;
384        assert!(escape.is_err());
385    }
386
387    #[tokio::test]
388    async fn db_schema_describes_csv_columns() {
389        let dir = tempfile::tempdir().unwrap();
390        std::fs::write(dir.path().join("people.csv"), "name,age\nada,36\n").unwrap();
391        let workspace = CandidateWorkspace::create(dir.path(), "n1", 0, "r1").unwrap();
392        let outcome = DbSchema
393            .apply(
394                &workspace,
395                &probe_call("db_schema", serde_json::json!({"path": "people.csv"})),
396                &entries()[1],
397            )
398            .await
399            .unwrap();
400        assert!(outcome.output.contains("name"));
401        assert!(outcome.output.contains("age"));
402    }
403}
perspt_agent/tools/families/db.rs

perspt_agent/tools/families/
db.rs
