use rhai::Dynamic; use rhai::Engine; use serde_json::{json, Value}; use sqlx::{PgPool}; use std::error::Error; use crate::services::state::AppState; use crate::services::utils; pub fn set_keyword(state: &AppState, engine: &mut Engine) { let db = state.db_custom.clone(); engine .register_custom_syntax(&["SET", "$expr$", ",", "$expr$", ",", "$expr$"], false, { let db = db.clone(); move |context, inputs| { let table_name = context.eval_expression_tree(&inputs[0])?; let filter = context.eval_expression_tree(&inputs[1])?; let updates = context.eval_expression_tree(&inputs[2])?; let binding = db.as_ref().unwrap(); // Use the current async context instead of creating a new runtime let binding2 = table_name.to_string(); let binding3 = filter.to_string(); let binding4 = updates.to_string(); let fut = execute_set(binding, &binding2, &binding3, &binding4); // Use tokio::task::block_in_place + tokio::runtime::Handle::current().block_on let result = tokio::task::block_in_place(|| tokio::runtime::Handle::current().block_on(fut)) .map_err(|e| format!("DB error: {}", e))?; if let Some(rows_affected) = result.get("rows_affected") { Ok(Dynamic::from(rows_affected.as_i64().unwrap_or(0))) } else { Err("No rows affected".into()) } } }) .unwrap(); } pub async fn execute_set( pool: &PgPool, table_str: &str, filter_str: &str, updates_str: &str, ) -> Result { println!( "Starting execute_set with table: {}, filter: {}, updates: {}", table_str, filter_str, updates_str ); // Parse the updates first to know how many parameters we'll have let (set_clause, update_params) = parse_updates(updates_str).map_err(|e| e.to_string())?; let update_params_count = update_params.len(); // Parse the filter condition with an offset for parameter indices let (where_clause, filter_params) = utils::parse_filter_with_offset(filter_str, update_params_count) .map_err(|e| e.to_string())?; // Combine all parameters (updates first, then filter) let mut params = update_params; params.extend(filter_params); let query = format!( "UPDATE {} SET {} WHERE {}", table_str, set_clause, where_clause ); println!("Executing query: {}", query); // Execute the update with all parameters let mut query_builder = sqlx::query(&query); for param in ¶ms { query_builder = query_builder.bind(param); } let result = query_builder .execute(pool) .await .map_err(|e| { eprintln!("SQL execution error: {}", e); e.to_string() })?; println!("Update successful, affected {} rows", result.rows_affected()); Ok(json!({ "command": "set", "table": table_str, "filter": filter_str, "updates": updates_str, "rows_affected": result.rows_affected() })) } // Helper function to parse the updates string into SQL SET clause and parameters fn parse_updates(updates_str: &str) -> Result<(String, Vec), Box> { let mut set_clauses = Vec::new(); let mut params = Vec::new(); // Split multiple updates by comma for (i, update) in updates_str.split(',').enumerate() { let parts: Vec<&str> = update.split('=').collect(); if parts.len() != 2 { return Err("Invalid update format. Expected 'KEY=VALUE'".into()); } let column = parts[0].trim(); let value = parts[1].trim(); // Validate column name to prevent SQL injection if !column.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { return Err("Invalid column name in update".into()); } set_clauses.push(format!("{} = ${}", column, i + 1)); params.push(value.to_string()); } Ok((set_clauses.join(", "), params)) }