From d659f48c248fdf6f3861f81516652dc2bcf2ddd9 Mon Sep 17 00:00:00 2001
From: Sage Griffin <sage@sagetheprogrammer.com>
Date: Tue, 4 Aug 2026 14:26:36 -0600
Subject: [PATCH] Add `MemoryToken::parse`

Currently if you want to manipulate a query immediately after parsing
it, you have to call `pg_raw_parse::parse`, then
`MemoryToken::make_unique` to perform a deep copy onto a new memory
context. We have places in `pgdog` which are using this pattern, and
immediately throwing away the unmodified AST/MemoryContext. We can skip
the extra copy/allocation by allowing you to parse within `make::try_owned`
directly.

We ended up never actually exposing the `warnings` field of
`ParseResult`, so I've gone ahead and removed it entirely. If we do ever
want to expose them, I'll either need to change the signature of `parse`
to return them as well as the stmts, or provide a new function that
returns both.

At this point `ParseResult` is a pointless type and can be removed, but
doing so is an API breaking change so I've left it in for now.
---
 src/lib.rs  | 47 ++++-------------------------------------------
 src/make.rs | 31 +++++++++++++++++++++++++++++--
 2 files changed, 33 insertions(+), 45 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 39d701f..e32f2a0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,5 @@
 #![cfg_attr(feature = "field_offset_assertions", feature(offset_of_enum))]
-use std::{ffi, fmt, ops, ptr};
+use std::{fmt, ops};
 
 pub mod const_val;
 mod deparse;
@@ -29,36 +29,14 @@ pub(crate) use node_ptr::{
 };
 
 pub fn parse(sql: &str) -> Result<ParseResult, error::Error> {
-    let mem = mem::MemoryContext::new(c"pg_raw_parse");
-    let cstring = ffi::CString::new(sql).map_err(error::Error::StatementContainedNul)?;
-    // SAFETY: we never panic within the provided block
-    let c_result = unsafe {
-        mem.within(|| {
-            raw::pg_query_raw_parse(
-                cstring.as_ptr(),
-                raw::PgQueryParseMode::PG_QUERY_PARSE_DEFAULT as _,
-            )
-        })
-    };
-    // Any warnings that were emitted during parsing went into a malloc'd
-    // buffer, so we need to construct this even if we're going to return Err
-    // to ensure that buffer is freed.
-    let warnings = Warnings {
-        stderr_buffer: ptr::NonNull::new(c_result.stderr_buffer),
-    };
-    match ptr::NonNull::new(c_result.error) {
-        Some(e) => Err(Error::from_pg_query_error(e)),
-        None => Ok(ParseResult {
-            _warnings: warnings,
-            tree: Owned::new(mem, c_result.tree.cast()),
-        }),
-    }
+    Ok(ParseResult {
+        tree: make::try_owned(|mem| mem.parse(sql))?,
+    })
 }
 
 pub type StmtList = list::CastNodeList<nodes::RawStmt>;
 
 pub struct ParseResult {
-    _warnings: Warnings,
     tree: Owned<StmtList>,
 }
 
@@ -85,20 +63,3 @@ impl fmt::Debug for ParseResult {
             .finish_non_exhaustive()
     }
 }
-
-struct Warnings {
-    stderr_buffer: Option<ptr::NonNull<ffi::c_char>>,
-}
-
-impl Drop for Warnings {
-    fn drop(&mut self) {
-        // tree was created with palloc, so is managed by postgres.
-        // stderr_buffer was malloc'd and must be freed
-        // SAFETY: libpg_query documents that the caller must free this.
-        unsafe {
-            if let Some(ptr) = self.stderr_buffer.take() {
-                libc::free(ptr.as_ptr() as _);
-            }
-        }
-    }
-}
diff --git a/src/make.rs b/src/make.rs
index db8f409..cf21fe5 100644
--- a/src/make.rs
+++ b/src/make.rs
@@ -2,11 +2,12 @@ use crate::list::{CastNodeList, NodeList};
 use crate::mem::MemoryContext;
 use crate::raw::{self, *};
 use crate::{
-    AsNodePtr, ConstValue, ConstructableNode, FromNodeMut, FromNodePtr, Node, Owned, nodes,
+    AsNodePtr, ConstValue, ConstructableNode, Error, FromNodeMut, FromNodePtr, Node, Owned, Result,
+    StmtList, nodes,
 };
 use generativity::Id;
 use std::any::type_name;
-use std::ffi::{c_char, c_int};
+use std::ffi::{CString, c_char, c_int};
 use std::marker::PhantomData;
 use std::ops::Deref;
 use std::ptr;
@@ -45,6 +46,32 @@ pub struct MemoryToken<'mem> {
 }
 
 impl<'mem> MemoryToken<'mem> {
+    /// Parse the given `sql` into a new AST on this memory context.
+    ///
+    /// This function can be used if you need to parse and then immediately
+    /// modify an AST, without copying it. If you only need to parse an AST,
+    /// use [`crate::parse`]
+    pub fn parse(self, sql: &str) -> Result<Unique<'mem, &'mem StmtList>> {
+        let cstring = CString::new(sql).map_err(Error::StatementContainedNul)?;
+        // SAFETY: we never panic within the provided block
+        let c_result = unsafe {
+            self.mem.within(|| {
+                raw::pg_query_raw_parse(
+                    cstring.as_ptr(),
+                    raw::PgQueryParseMode::PG_QUERY_PARSE_DEFAULT as _,
+                )
+            })
+        };
+        if !c_result.stderr_buffer.is_null() {
+            // SAFETY: libpg_query documents that the caller must free this.
+            unsafe { libc::free(c_result.stderr_buffer as _) };
+        }
+        match ptr::NonNull::new(c_result.error) {
+            Some(e) => Err(Error::from_pg_query_error(e)),
+            None => Ok(Unique(c_result.tree.cast(), self.id, PhantomData)),
+        }
+    }
+
     pub fn make_a_const(self, val: ConstValue<'_>) -> Unique<'mem, &'mem nodes::A_Const> {
         let mut node = self.make_node::<nodes::A_Const>();
         node.as_mut().set_isnull(false);