diff --git a/Cargo.lock b/Cargo.lock index dd9a35a..c11528c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4028,7 +4028,7 @@ dependencies = [ "serde_json", "signal-hook", "soa_derive", - "sqlparser 0.50.0", + "sqlparser 0.51.0", "sqlx", "strum 0.26.3", "supabase-wrappers", @@ -5405,9 +5405,7 @@ dependencies = [ [[package]] name = "sqlparser" -version = "0.50.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2e5b515a2bd5168426033e9efbfd05500114833916f1d5c268f938b4ee130ac" +version = "0.51.0" dependencies = [ "log", ] diff --git a/Cargo.toml b/Cargo.toml index 020fd67..2363926 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,7 +29,7 @@ pgrx = "0.12.5" serde = "1.0.210" serde_json = "1.0.128" signal-hook = "0.3.17" -sqlparser = "0.50.0" +sqlparser = { path = "/home/kyss/labs/sqlparser-rs" } strum = { version = "0.26.3", features = ["derive"] } supabase-wrappers = { git = "https://github.com/paradedb/wrappers.git", default-features = false, rev = "c32abb7" } thiserror = "1.0.63" diff --git a/src/duckdb/connection.rs b/src/duckdb/connection.rs index b420916..a40ac0f 100644 --- a/src/duckdb/connection.rs +++ b/src/duckdb/connection.rs @@ -249,3 +249,21 @@ pub fn set_search_path(search_path: Vec) -> Result<()> { Ok(()) } + +pub fn execute_explain(query: &str) -> Result { + let conn = unsafe { &*get_global_connection().get() }; + let mut stmt = conn.prepare(query)?; + let rows = stmt.query_row([], |row| { + let mut r = vec![]; + + let mut col_index = 1; + while let Ok(value) = row.get::<_, String>(col_index) { + r.push(value); + col_index += 1; + } + + Ok(r) + })?; + + Ok(rows.join("")) +} diff --git a/src/hooks/utility.rs b/src/hooks/utility.rs index d53dbd4..e58ffe3 100644 --- a/src/hooks/utility.rs +++ b/src/hooks/utility.rs @@ -154,6 +154,7 @@ fn parse_query_from_utility_stmt(query_string: &core::ffi::CStr) -> Result Ok(statement.to_string()), _ => bail!("unexpected utility statement: {}", query_string), } diff --git a/src/hooks/utility/explain.rs b/src/hooks/utility/explain.rs index bdfca82..96e6463 100644 --- a/src/hooks/utility/explain.rs +++ b/src/hooks/utility/explain.rs @@ -15,13 +15,26 @@ // You should have received a copy of the GNU Affero General Public License // along with this program. If not, see . -use std::ffi::CString; +use std::ffi::{CStr, CString}; +use std::time::Instant; use anyhow::Result; use pgrx::{error, pg_sys}; use super::parse_query_from_utility_stmt; -use crate::hooks::query::{get_query_relations, is_duckdb_query}; +use crate::{ + duckdb::connection, + hooks::query::{get_query_relations, is_duckdb_query, set_search_path_by_pg}, +}; + +enum Style { + Postgres, + Duckdb, +} +struct ExplainState { + analyze: bool, + style: Style, +} pub fn explain_query( query_string: &core::ffi::CStr, @@ -37,9 +50,34 @@ pub fn explain_query( return Ok(true); } - if unsafe { !(*stmt).options.is_null() } { - error!("the EXPLAIN options provided are not supported for DuckDB pushdown queries."); - } + let state = parse_explain_options(unsafe { (*stmt).options }); + let query = parse_query_from_utility_stmt(query_string)?; + + let output = match state.style { + Style::Postgres => { + let mut output = format!("DuckDB Scan: {}\n", query); + if state.analyze { + let start_time = Instant::now(); + set_search_path_by_pg()?; + connection::execute(&query, [])?; + let duration = start_time.elapsed(); + output += &format!( + "Execution Time: {:.3} ms\n", + duration.as_micros() as f64 / 1_000.0 + ); + } + output + } + Style::Duckdb => { + set_search_path_by_pg()?; + let explain_query = if state.analyze { + format!("EXPLAIN ANALYZE {query}") + } else { + format!("EXPLAIN {query}") + }; + connection::execute_explain(&explain_query)? + } + }; unsafe { let tstate = pg_sys::begin_tup_output_tupdesc( @@ -47,15 +85,71 @@ pub fn explain_query( pg_sys::ExplainResultDesc(stmt), &pg_sys::TTSOpsVirtual, ); - let query = format!( - "DuckDB Scan: {}", - parse_query_from_utility_stmt(query_string)? - ); - let query_c_str = CString::new(query)?; - pg_sys::do_text_output_multiline(tstate, query_c_str.as_ptr()); + let output_cstr = CString::new(output)?; + + pg_sys::do_text_output_multiline(tstate, output_cstr.as_ptr()); pg_sys::end_tup_output(tstate); } Ok(false) } + +fn parse_explain_options(options: *const pg_sys::List) -> ExplainState { + let mut explain_state = ExplainState { + analyze: false, + style: Style::Postgres, + }; + + if options.is_null() { + return explain_state; + } + + unsafe { + let elements = (*options).elements; + + for i in 0..(*options).length as isize { + let opt = (*elements.offset(i)).ptr_value as *mut pg_sys::DefElem; + + let opt_name = match CStr::from_ptr((*opt).defname).to_str() { + Ok(opt) => opt, + Err(e) => { + error!("failed to parse EXPLAIN option name: {e}"); + } + }; + match opt_name { + "analyze" => { + explain_state.analyze = pg_sys::defGetBoolean(opt); + } + "style" => { + let style = match CStr::from_ptr(pg_sys::defGetString(opt)).to_str() { + Ok(style) => style, + + Err(e) => { + error!("failed to parse STYLE option: {e}"); + } + }; + + explain_state.style = match parse_explain_style(style) { + Some(s) => s, + None => { + error!("unrecognized STYLE option: {style}") + } + }; + } + _ => error!("unrecognized EXPLAIN option \"{opt_name}\""), + } + } + } + + explain_state +} + +fn parse_explain_style(style: &str) -> Option