diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs index 74e7401..9aa8d19 100644 --- a/src/handlers/mod.rs +++ b/src/handlers/mod.rs @@ -6,3 +6,4 @@ pub mod crud; pub mod ddl; pub mod metadata; pub mod query; +pub mod routines; diff --git a/src/handlers/routines.rs b/src/handlers/routines.rs new file mode 100644 index 0000000..7a2ac74 --- /dev/null +++ b/src/handlers/routines.rs @@ -0,0 +1,207 @@ +//! PostgreSQL-dialect SQL builders and RPC handlers for stored-routine +//! management: `build_routine_call_sql`, `routine_create_template`, +//! `drop_routine`. +//! +//! Mirrors the built-in driver exactly +//! (`src-tauri/src/drivers/postgres/routines.rs` for the pure string +//! builders, `mod.rs`'s `drop_routine` for the live identity-signature +//! lookup) so both drivers produce byte-identical SQL for the same inputs. +//! `get_routine_edit_script` is intentionally not implemented here: both the +//! builtin and the host's plugin-bridge fallback resolve it to +//! `get_routine_definition`, which this plugin already implements. + +use serde_json::{json, Value}; + +use crate::client; +use crate::models::{inner_params, ConnectionParams, RoutineCallArg}; +use crate::rpc::{error_response, ok_response}; +use crate::utils::identifiers::{qualified, quote_identifier}; + +pub async fn build_routine_call_sql(id: Value, params: &Value) -> Value { + let routine_name = params + .get("routine_name") + .and_then(Value::as_str) + .unwrap_or(""); + let routine_type = params + .get("routine_type") + .and_then(Value::as_str) + .unwrap_or("FUNCTION"); + let schema = params + .get("schema") + .and_then(Value::as_str) + .unwrap_or("public"); + let args: Vec = params + .get("args") + .and_then(|v| serde_json::from_value(v.clone()).ok()) + .unwrap_or_default(); + + let sql = routine_call_sql(routine_name, routine_type, &args, schema); + ok_response(id, json!(sql)) +} + +pub async fn routine_create_template(id: Value, params: &Value) -> Value { + let routine_type = params + .get("routine_type") + .and_then(Value::as_str) + .unwrap_or("FUNCTION"); + let schema = params + .get("schema") + .and_then(Value::as_str) + .unwrap_or("public"); + + ok_response(id, json!(routine_template(routine_type, schema))) +} + +pub async fn drop_routine(id: Value, params: &Value) -> Value { + let conn_params = ConnectionParams::from_value(inner_params(params)); + let routine_name = params + .get("routine_name") + .and_then(Value::as_str) + .unwrap_or(""); + let routine_type = params + .get("routine_type") + .and_then(Value::as_str) + .unwrap_or("FUNCTION"); + let schema = params + .get("schema") + .and_then(Value::as_str) + .unwrap_or("public"); + + match exec_drop_routine(&conn_params, routine_name, routine_type, schema).await { + Ok(()) => ok_response(id, Value::Null), + Err(e) => error_response(id, -32603, &e), + } +} + +/// Drops a routine, resolving its exact identity signature first: PostgreSQL +/// identifies routines by name *and* argument types, so a bare +/// `DROP FUNCTION name` fails as soon as overloads exist. +async fn exec_drop_routine( + conn_params: &ConnectionParams, + routine_name: &str, + routine_type: &str, + schema: &str, +) -> Result<(), String> { + let query = r#" + SELECT pg_get_function_identity_arguments(p.oid) AS args + FROM pg_proc p + JOIN pg_namespace n ON p.pronamespace = n.oid + WHERE n.nspname = $1 AND p.proname = $2 + "#; + let rows = client::query_rows(conn_params, query, &[&schema, &routine_name]).await?; + + match rows.len() { + 0 => Err(format!( + "Routine '{}' not found in schema '{}'", + routine_name, schema + )), + 1 => { + let identity_args: String = rows[0].try_get("args").unwrap_or_default(); + let sql = drop_routine_sql(routine_name, routine_type, &identity_args, schema); + client::execute_typed(conn_params, &sql, &[]) + .await + .map(|_| ()) + } + n => Err(format!( + "Routine '{}' has {} overloads; drop it manually specifying the argument types", + routine_name, n + )), + } +} + +/// Builds the invocation script. Functions go through `SELECT * FROM` so +/// both scalar and set-returning functions come back as a result set; their +/// pure `OUT` parameters are NOT part of the call signature in PostgreSQL, so +/// they are excluded from the argument list (passing them raises +/// `function ... does not exist`). Procedures use `CALL`; there OUT +/// parameters ARE required in the argument list and are rendered as `NULL` +/// placeholders, with INOUT values echoed back by the server as the +/// procedure's result row. +fn routine_call_sql( + routine_name: &str, + routine_type: &str, + args: &[RoutineCallArg], + schema: &str, +) -> String { + let name = qualified(schema, routine_name); + let is_function = routine_type.eq_ignore_ascii_case("FUNCTION"); + let rendered: Vec = args + .iter() + .filter(|arg| !(is_function && arg.mode.eq_ignore_ascii_case("OUT"))) + .map(render_sql_literal) + .collect(); + let arg_list = rendered.join(", "); + if is_function { + format!("SELECT * FROM {}({});", name, arg_list) + } else { + format!("CALL {}({});", name, arg_list) + } +} + +fn render_sql_literal(arg: &RoutineCallArg) -> String { + match &arg.value { + None => "NULL".to_string(), + Some(v) if arg.is_raw => v.clone(), + Some(v) => format!("'{}'", v.replace('\'', "''")), + } +} + +/// Starter script for a new routine. `CREATE OR REPLACE` keeps the script +/// re-runnable while iterating on the body. +fn routine_template(routine_type: &str, schema: &str) -> String { + let prefix = if schema.is_empty() { + String::new() + } else { + format!("{}.", quote_identifier(schema)) + }; + if routine_type.eq_ignore_ascii_case("FUNCTION") { + format!( + r#"CREATE OR REPLACE FUNCTION {prefix}my_function(p_value integer) +RETURNS integer +LANGUAGE plpgsql +AS $$ +BEGIN + RETURN p_value; +END; +$$; +"# + ) + } else { + format!( + r#"CREATE OR REPLACE PROCEDURE {prefix}my_procedure(p_value integer) +LANGUAGE plpgsql +AS $$ +BEGIN + RAISE NOTICE 'value: %', p_value; +END; +$$; +"# + ) + } +} + +/// `DROP` statement for a routine identified by its exact signature (the +/// output of `pg_get_function_identity_arguments`), which is how PostgreSQL +/// disambiguates overloads. +fn drop_routine_sql( + routine_name: &str, + routine_type: &str, + identity_args: &str, + schema: &str, +) -> String { + let keyword = if routine_type.eq_ignore_ascii_case("PROCEDURE") { + "PROCEDURE" + } else { + "FUNCTION" + }; + format!( + "DROP {} {}({})", + keyword, + qualified(schema, routine_name), + identity_args + ) +} + +#[cfg(test)] +#[path = "routines_tests.rs"] +mod routines_tests; diff --git a/src/handlers/routines_tests.rs b/src/handlers/routines_tests.rs new file mode 100644 index 0000000..fe26541 --- /dev/null +++ b/src/handlers/routines_tests.rs @@ -0,0 +1,145 @@ +//! Unit tests for `handlers/routines.rs`'s pure SQL builders. Sibling test +//! file per repo convention (`#[cfg(test)] #[path = ...] mod routines_tests;` +//! in `routines.rs`) — mirrors the builtin driver's own +//! `routine_management` test module +//! (`src-tauri/src/drivers/postgres/tests.rs`) so the two stay provably in +//! sync on expected output. + +use super::{drop_routine_sql, render_sql_literal, routine_call_sql, routine_template}; +use crate::models::RoutineCallArg; + +fn arg(name: &str, mode: &str, value: Option<&str>, is_raw: bool) -> RoutineCallArg { + RoutineCallArg { + name: name.to_string(), + mode: mode.to_string(), + value: value.map(|v| v.to_string()), + is_raw, + } +} + +#[test] +fn function_call_uses_select_star_from() { + let sql = routine_call_sql( + "fn_report", + "FUNCTION", + &[arg("p_year", "IN", Some("2026"), true)], + "public", + ); + assert_eq!(sql, "SELECT * FROM \"public\".\"fn_report\"(2026);"); +} + +#[test] +fn function_call_excludes_out_params() { + // PostgreSQL functions do not accept pure OUT parameters in the call + // signature; only IN/INOUT are passed. + let sql = routine_call_sql( + "fn_split", + "FUNCTION", + &[ + arg("p_in", "IN", Some("5"), true), + arg("p_lo", "OUT", None, false), + arg("p_hi", "OUT", None, false), + ], + "public", + ); + assert_eq!(sql, "SELECT * FROM \"public\".\"fn_split\"(5);"); +} + +#[test] +fn function_call_keeps_inout_params() { + let sql = routine_call_sql( + "fn_adjust", + "FUNCTION", + &[ + arg("p_val", "INOUT", Some("10"), true), + arg("p_out", "OUT", None, false), + ], + "public", + ); + assert_eq!(sql, "SELECT * FROM \"public\".\"fn_adjust\"(10);"); +} + +#[test] +fn procedure_call_renders_out_params_as_null() { + let sql = routine_call_sql( + "sp_test", + "PROCEDURE", + &[ + arg("p_in", "IN", Some("it's"), false), + arg("p_out", "OUT", None, false), + ], + "public", + ); + assert_eq!(sql, "CALL \"public\".\"sp_test\"('it''s', NULL);"); +} + +#[test] +fn procedure_call_keeps_out_params_in_argument_list() { + // Unlike functions, procedures require every OUT parameter in the call + // signature — this is the "OUT args ARE required" half of the contract. + let sql = routine_call_sql( + "sp_split", + "PROCEDURE", + &[ + arg("p_in", "IN", Some("5"), true), + arg("p_out", "OUT", None, false), + ], + "public", + ); + assert_eq!(sql, "CALL \"public\".\"sp_split\"(5, NULL);"); +} + +#[test] +fn render_sql_literal_quotes_and_escapes_non_raw_values() { + assert_eq!( + render_sql_literal(&arg("p", "IN", Some("it's"), false)), + "'it''s'" + ); +} + +#[test] +fn render_sql_literal_passes_raw_values_verbatim() { + assert_eq!( + render_sql_literal(&arg("p", "IN", Some("now()"), true)), + "now()" + ); +} + +#[test] +fn render_sql_literal_renders_missing_value_as_null() { + assert_eq!(render_sql_literal(&arg("p", "OUT", None, false)), "NULL"); +} + +#[test] +fn create_templates_are_schema_qualified_or_replace() { + let tpl = routine_template("FUNCTION", "app"); + assert!( + tpl.starts_with("CREATE OR REPLACE FUNCTION \"app\"."), + "{tpl}" + ); + let tpl = routine_template("PROCEDURE", ""); + assert!( + tpl.starts_with("CREATE OR REPLACE PROCEDURE my_procedure"), + "{tpl}" + ); +} + +#[test] +fn function_template_is_valid_plpgsql_with_dollar_quoting() { + let tpl = routine_template("FUNCTION", "public"); + assert!(tpl.contains("LANGUAGE plpgsql")); + assert!(tpl.contains("AS $$")); + assert!(tpl.contains("RETURNS integer")); +} + +#[test] +fn drop_sql_includes_identity_signature() { + assert_eq!( + drop_routine_sql("fn_add", "FUNCTION", "integer, integer", "public"), + "DROP FUNCTION \"public\".\"fn_add\"(integer, integer)" + ); + assert_eq!( + drop_routine_sql("sp", "PROCEDURE", "", "public"), + "DROP PROCEDURE \"public\".\"sp\"()" + ); +} diff --git a/src/models.rs b/src/models.rs index 0b7eedb..2f3ab82 100644 --- a/src/models.rs +++ b/src/models.rs @@ -69,3 +69,16 @@ pub struct ColumnDefinition { pub is_auto_increment: bool, pub default_value: Option, } + +/// Mirrors `crate::models::RoutineCallArg` on the host — one argument of a +/// stored-routine invocation built by `build_routine_call_sql`. +#[derive(Debug, Clone, Deserialize)] +pub struct RoutineCallArg { + pub name: String, + /// "IN", "OUT", or "INOUT". + pub mode: String, + #[serde(default)] + pub value: Option, + #[serde(default)] + pub is_raw: bool, +} diff --git a/src/rpc.rs b/src/rpc.rs index 1b24690..4b6e8df 100644 --- a/src/rpc.rs +++ b/src/rpc.rs @@ -50,6 +50,9 @@ pub async fn handle_line(line: &str) -> Value { "get_routines" => handlers::metadata::get_routines(id, ¶ms).await, "get_routine_parameters" => handlers::metadata::get_routine_parameters(id, ¶ms).await, "get_routine_definition" => handlers::metadata::get_routine_definition(id, ¶ms).await, + "build_routine_call_sql" => handlers::routines::build_routine_call_sql(id, ¶ms).await, + "routine_create_template" => handlers::routines::routine_create_template(id, ¶ms).await, + "drop_routine" => handlers::routines::drop_routine(id, ¶ms).await, "get_triggers" => handlers::metadata::get_triggers(id, ¶ms).await, "get_trigger_definition" => handlers::metadata::get_trigger_definition(id, ¶ms).await, "get_schema_snapshot" => handlers::metadata::get_schema_snapshot(id, ¶ms).await,