From 3871739513387f2473fbc961840de5951e7ed677 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sat, 27 Jun 2026 14:34:28 -0600 Subject: [PATCH] Add the ability to deparse a statement This wraps `deparesRawStmt` from libpg_query in a safe Rust wrapper. I was a bit surprised to find that PG itself doesn't contain any deparsing functions that are usable for us, so we will need to continue depending on libpg_query for the foreseeable future rather than PG itself, unless we want to write our own deparser (which actually may be worth doing). In order to make writing the tests easier, I needed the helper methods on `List` to be available on cast lists. I went ahead and copied all of them over while I was there. --- build.rs | 2 ++ src/deparse.rs | 64 +++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 16 +++++++------ src/list.rs | 35 +++++++++++++++++++++++++++ src/pg_error.rs | 6 +++++ wrapper.h | 12 ++++++++++ 6 files changed, 128 insertions(+), 7 deletions(-) create mode 100644 src/deparse.rs diff --git a/build.rs b/build.rs index 3ddb451..8e6de6b 100644 --- a/build.rs +++ b/build.rs @@ -84,6 +84,8 @@ fn main() { bindgen::Abi::CUnwind, "wrapped_raw_expression_tree_walker_impl", ) + .allowlist_item("StringInfo") + .allowlist_item("wrapped_raw_deparse") .wrap_static_fns(true) .wrap_static_fns_path(out_dir.join("wrap_static_fns")); for struct_name in &node_structs { diff --git a/src/deparse.rs b/src/deparse.rs new file mode 100644 index 0000000..cf7af08 --- /dev/null +++ b/src/deparse.rs @@ -0,0 +1,64 @@ +use crate::mem::MemoryContext; +use crate::pg_error::PgError; +use crate::{Result, nodes, raw}; +use std::ptr::{self, NonNull}; + +pub struct DeparseResult { + _mem: MemoryContext, + string_info: raw::StringInfo, +} + +impl DeparseResult { + pub fn deparse(stmt: &nodes::RawStmt) -> Result { + let mem = MemoryContext::new(c"pg_raw_deparse"); + let mut err = ptr::null_mut(); + // SAFETY: We never panic + let string_info = unsafe { + mem.within(|| raw::wrapped_raw_deparse((&raw const *stmt).cast_mut(), &mut err)) + }; + + match NonNull::new(err) { + Some(err) => Err(PgError::new(mem, err).into()), + None => Ok(Self { + _mem: mem, + string_info, + }), + } + } + + pub fn as_str(&self) -> &str { + // SAFETY: We're always allocating a valid pointer + let info = unsafe { &*self.string_info }; + // SAFETY: We're always allocating a valid pointer + let bytes = + unsafe { std::slice::from_raw_parts(info.data.cast_const().cast(), info.len as usize) }; + std::str::from_utf8(bytes).expect("PG always returns valid UTF-8") + } +} + +pub fn deparse(stmt: &nodes::RawStmt) -> Result { + DeparseResult::deparse(stmt) +} + +#[test] +fn test_deparse() { + fn run_test(query: &str) { + let result = crate::parse(query).unwrap(); + let stmt = result.raw_stmts().first().unwrap(); + let deparsed = deparse(stmt).unwrap(); + + assert_eq!(query, deparsed.as_str()); + } + + run_test("SELECT 1"); + run_test("SELECT id, name, email FROM users WHERE users.id = 1"); + run_test( + "UPDATE users SET name = 'Sage', email = 'sage@pgdog.dev' WHERE users.id = 1 AND NOT EXISTS (SELECT 1 FROM users WHERE email = 'sage@pgdog.dev')", + ); + run_test( + "INSERT INTO users (name, email) VALUES ('Sage', 'sage@pgdog.dev') ON CONFLICT DO NOTHING", + ); + run_test("DELETE FROM users WHERE id = 1"); + run_test("SET my_config TO 1"); + run_test("TRUNCATE users"); +} diff --git a/src/lib.rs b/src/lib.rs index d9e4758..055898e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ use std::{ffi, fmt, ptr}; pub mod const_val; +mod deparse; pub mod error; pub mod list; mod mem; @@ -13,6 +14,7 @@ pub mod raw; pub mod walk; pub use crate::const_val::ConstValue; +pub use crate::deparse::{DeparseResult, deparse}; pub use crate::error::{Error, Result}; pub use crate::node_enum::Node; @@ -53,18 +55,18 @@ unsafe impl Send for ParseResult {} unsafe impl Sync for ParseResult {} impl ParseResult { - /// Returns the list of statements received, panics if the list was a - /// type other than Node + /// Returns the statements that were parsed pub fn stmts(&self) -> impl Iterator> { + self.raw_stmts().into_iter().map(|s| s.stmt()) + } + + /// Returns the raw statements that were parsed + pub fn raw_stmts(&self) -> &list::CastNodeList<&nodes::RawStmt> { // SAFETY: The memory context of the tree is guaranteed to outlive // the lifetime of self. We are returning a lifetime shorter than self. unsafe { Node::from_ptr(self.tree.tree.cast()) } .expect_node_list() - .into_iter() - .map(|n| match n { - Node::RawStmt(stmt) => stmt.stmt(), - n => n, - }) + .cast() } } diff --git a/src/list.rs b/src/list.rs index 7e044e8..5c98b11 100644 --- a/src/list.rs +++ b/src/list.rs @@ -65,6 +65,11 @@ impl NodeList { pub fn is_empty(&self) -> bool { self.len() == 0 } + + #[inline] + pub fn first(&self) -> Option<<&Self as IntoIterator>::Item> { + self.into_iter().next() + } } impl<'a> IntoIterator for &'a NodeList { @@ -191,6 +196,36 @@ pub struct CastNodeList { _marker: PhantomData, } +impl CastNodeList { + #[inline] + pub fn iter<'a>(&'a self) -> <&'a Self as IntoIterator>::IntoIter + where + &'a Self: IntoIterator, + { + self.into_iter() + } + + #[inline] + #[must_use] + pub fn len(&self) -> usize { + self.list.length as usize + } + + #[inline] + #[must_use] + pub fn is_empty(&self) -> bool { + self.list.len() == 0 + } + + #[inline] + pub fn first<'a>(&'a self) -> Option<<&'a Self as IntoIterator>::Item> + where + &'a Self: IntoIterator, + { + self.into_iter().next() + } +} + #[allow(clippy::unnecessary_operation, clippy::identity_op)] const _: () = { use crate::nodes::String; diff --git a/src/pg_error.rs b/src/pg_error.rs index fbf6dc7..b4b82c1 100644 --- a/src/pg_error.rs +++ b/src/pg_error.rs @@ -16,6 +16,12 @@ unsafe impl Send for PgError {} unsafe impl Sync for PgError {} impl PgError { + pub(crate) fn new(mem: MemoryContext, error_data: NonNull) -> Self { + Self { + _mem: mem, + error_data: ErrorData(error_data), + } + } /// Returns None if the raw error is NULL pub(crate) fn from_raw(raw: raw::Error) -> Option { NonNull::new(raw.error_data).map(|error_data| { diff --git a/wrapper.h b/wrapper.h index 76808f0..9f7f095 100644 --- a/wrapper.h +++ b/wrapper.h @@ -21,3 +21,15 @@ static inline bool wrapped_raw_expression_tree_walker_impl(Node *n, tree_walker_ PG_END_TRY(); return result; } + +static inline StringInfo wrapped_raw_deparse(RawStmt *stmt, ErrorData **error) { + StringInfo str; + PG_TRY(); + str = makeStringInfo(); + deparseRawStmt(str, stmt); + PG_CATCH(); + *error = CopyErrorData(); + FlushErrorState(); + PG_END_TRY(); + return str; +}