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; +}