diff --git a/server/src/steel/server/execution.rs b/server/src/steel/server/execution.rs index 6a45b45..4cc14c9 100644 --- a/server/src/steel/server/execution.rs +++ b/server/src/steel/server/execution.rs @@ -52,6 +52,15 @@ pub fn execute_script( } }); + // SQL query registration + vm.register_fn("steel_query_sql", { + let ctx = context.clone(); + move |query: String| { + ctx.steel_query_sql(&query) + .map_err(|e| e.to_string()) + } + }); + // Execute script and process results let results = vm.compile_and_run_raw_program(script) .map_err(|e| ExecutionError::RuntimeError(e.to_string()))?; diff --git a/server/src/steel/server/functions.rs b/server/src/steel/server/functions.rs index adceef1..17397a9 100644 --- a/server/src/steel/server/functions.rs +++ b/server/src/steel/server/functions.rs @@ -91,4 +91,40 @@ impl SteelContext { Err(SteelVal::StringV("Expected comma-separated string".into())) } } + + pub fn steel_query_sql(&self, query: &str) -> Result { + // Validate query is read-only + if !is_read_only_query(query) { + return Err(SteelVal::StringV( + "Only SELECT queries are allowed".into() + )); + } + + let pool = self.db_pool.clone(); + let result = tokio::runtime::Runtime::new().unwrap().block_on(async { + // Execute and get first column of all rows as strings + let rows = sqlx::query(query) + .fetch_all(&*pool) + .await + .map_err(|e| SteelVal::StringV(e.to_string().into()))?; + + let mut results = Vec::new(); + for row in rows { + let val: String = row.try_get(0) + .map_err(|e| SteelVal::StringV(e.to_string().into()))?; + results.push(val); + } + + Ok(results.join(",")) + }); + + result.map(|s| SteelVal::StringV(s.into())) + } +} + +fn is_read_only_query(query: &str) -> bool { + let query = query.trim_start().to_uppercase(); + query.starts_with("SELECT") || + query.starts_with("SHOW") || + query.starts_with("EXPLAIN") }