feat: Enhance explain (#146)

* feat: enhance explain

* test: add duckdb style explain test

Cargo.toml

@@ -28,7 +28,7 @@ duckdb = { git = "https://github.com/paradedb/duckdb-rs.git", features = [
 pgrx = "0.12.7"
 serde_json = "1.0.128"
 signal-hook = "0.3.17"
-sqlparser = "0.50.0"
+sqlparser = "0.52.0"
 strum = { version = "0.26.3", features = ["derive"] }
 supabase-wrappers = { git = "https://github.com/paradedb/wrappers.git", default-features = false, rev = "8aef4a6" }
 thiserror = "1.0.63"

src/duckdb/connection.rs

@@ -66,6 +66,7 @@ pub fn get_global_connection() -> &'static UnsafeCell<Connection> {
     INIT.call_once(|| {
         init_globals();
     });
+    #[allow(static_mut_refs)]
     unsafe {
         GLOBAL_CONNECTION
             .as_ref()
@@ -77,6 +78,7 @@ fn get_global_statement() -> &'static UnsafeCell<Option<Statement<'static>>> {
     INIT.call_once(|| {
         init_globals();
     });
+    #[allow(static_mut_refs)]
     unsafe {
         GLOBAL_STATEMENT
             .as_ref()
@@ -88,7 +90,10 @@ fn get_global_arrow() -> &'static UnsafeCell<Option<duckdb::Arrow<'static>>> {
     INIT.call_once(|| {
         init_globals();
     });
-    unsafe { GLOBAL_ARROW.as_ref().expect("Arrow not initialized") }
+    #[allow(static_mut_refs)]
+    unsafe {
+        GLOBAL_ARROW.as_ref().expect("Arrow not initialized")
+    }
 }
 
 pub fn create_csv_view(
@@ -249,3 +254,21 @@ pub fn set_search_path(search_path: Vec<String>) -> Result<()> {
 
     Ok(())
 }
+
+pub fn execute_explain(query: &str) -> Result<String> {
+    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(""))
+}

src/duckdb/json.rs

@@ -98,10 +98,10 @@ fn extract_option(
     table_options: &HashMap<String, String>,
     quote: bool,
 ) -> Option<String> {
-    return table_options.get(option.as_ref()).map(|res| match quote {
+    table_options.get(option.as_ref()).map(|res| match quote {
         true => format!("{option} = '{res}'"),
         false => format!("{option} = {res}"),
-    });
+    })
 }
 
 #[cfg(test)]

src/hooks/utility.rs

@@ -163,13 +163,7 @@ fn parse_query_from_utility_stmt(query_string: &core::ffi::CStr) -> Result<Strin
 
     debug_assert!(utility.len() == 1);
     match &utility[0] {
-        Statement::Explain {
-            describe_alias: _,
-            analyze: _,
-            verbose: _,
-            statement,
-            format: _,
-        } => Ok(statement.to_string()),
+        Statement::Explain { statement, .. } => Ok(statement.to_string()),
         _ => bail!("unexpected utility statement: {}", query_string),
     }
 }

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 <http://www.gnu.org/licenses/>.
 
-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,25 +50,106 @@ 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(
             dest,
             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<Style> {
+    match style {
+        "pg" => Some(Style::Postgres),
+        "postgres" => Some(Style::Postgres),
+        "duckdb" => Some(Style::Duckdb),
+        _ => None,
+    }
+}