use serde_json::{Value, json};
use saya_agent::ToolError;
use super::DatabaseTools;
use crate::connection::ConnectionEntry;
impl DatabaseTools {
pub(super) async fn result_shape(
&self,
entry: &ConnectionEntry,
sql: &str,
) -> Result<Value, ToolError> {
self.detect_and_record_overrides(sql, entry.dialect);
let result = crate::agent::state_tools::query(
entry.connector.as_ref(),
sql,
self.max_rows,
self.state_db.as_ref(),
entry.profile_id.as_deref(),
)
.await?;
Ok(shape_of(&result))
}
}
fn shape_of(result: &Value) -> Value {
let row_count = result.get("row_count").and_then(Value::as_u64).unwrap_or(0);
let truncated = result
.get("truncated")
.and_then(Value::as_bool)
.unwrap_or(false);
let first_row = result
.get("rows")
.and_then(Value::as_array)
.and_then(|rows| rows.first())
.and_then(Value::as_array);
let columns = result
.get("columns")
.and_then(Value::as_array)
.map(|names| {
names
.iter()
.enumerate()
.map(|(index, name)| {
let ty = first_row
.and_then(|cells| cells.get(index))
.map(value_type_label)
.unwrap_or(Value::Null);
json!({ "name": name, "type": ty })
})
.collect::<Vec<_>>()
})
.unwrap_or_default();
json!({
"row_count": row_count,
"truncated": truncated,
"columns": columns,
})
}
fn value_type_label(value: &Value) -> Value {
let label = match value {
Value::Null => "NULL",
Value::Bool(_) => "BOOLEAN",
Value::Number(n) if n.is_i64() || n.is_u64() => "INTEGER",
Value::Number(_) => "REAL",
Value::String(_) => "TEXT",
Value::Array(_) | Value::Object(_) => "JSON",
};
Value::String(label.into())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn shape_of_infers_types_from_the_first_row_and_omits_values() {
let result = json!({
"columns": ["customer", "age", "balance", "active", "meta"],
"rows": [["Alice", 30, 99.5, true, {"k": 1}]],
"row_count": 1,
"truncated": false,
"executed_sql": "SELECT 1"
});
let shape = shape_of(&result);
assert_eq!(shape["row_count"], 1);
assert_eq!(shape["truncated"], false);
assert_eq!(shape["columns"][0]["name"], "customer");
assert_eq!(shape["columns"][0]["type"], "TEXT");
assert_eq!(shape["columns"][1]["type"], "INTEGER");
assert_eq!(shape["columns"][2]["type"], "REAL");
assert_eq!(shape["columns"][3]["type"], "BOOLEAN");
assert_eq!(shape["columns"][4]["type"], "JSON");
let serialized = serde_json::to_string(&shape).unwrap();
assert!(!serialized.contains("Alice"), "no cell value may escape");
}
#[test]
fn shape_of_keeps_column_names_with_null_types_for_an_empty_result() {
let result = json!({
"columns": ["customer", "age"],
"rows": [],
"row_count": 0,
"truncated": false,
"executed_sql": "SELECT 1"
});
let shape = shape_of(&result);
assert_eq!(shape["row_count"], 0);
let columns = shape["columns"].as_array().unwrap();
assert_eq!(columns.len(), 2);
assert_eq!(columns[0]["name"], "customer");
assert_eq!(columns[0]["type"], Value::Null);
assert_eq!(columns[1]["name"], "age");
assert_eq!(columns[1]["type"], Value::Null);
}
}