feat(pgls_query): add PLpgSQL protobuf parsing support

.gitmodules

@@ -1,4 +1,4 @@
 [submodule "crates/pgls_query/vendor/libpg_query"]
 	path = crates/pgls_query/vendor/libpg_query
-	url = https://github.com/pganalyze/libpg_query.git
-	branch = 17-latest
+	url = https://github.com/psteinroe/libpg_query.git
+	branch = feat/plpgsql-protobuf

crates/pgls_query/build.rs

@@ -100,7 +100,9 @@ fn run_bindgen(
         .allowlist_function("pg_query_split_with_parser")
         .allowlist_function("pg_query_split_with_scanner")
         .allowlist_function("pg_query_parse_plpgsql")
+        .allowlist_function("pg_query_parse_plpgsql_protobuf")
         .allowlist_function("pg_query_free_protobuf_parse_result")
+        .allowlist_function("pg_query_free_plpgsql_protobuf_parse_result")
         .allowlist_function("pg_query_free_scan_result")
         .allowlist_function("pg_query_free_deparse_result")
         .allowlist_function("pg_query_free_normalize_result")
@@ -117,6 +119,7 @@ fn run_bindgen(
         .allowlist_type("PgQueryFingerprintResult")
         .allowlist_type("PgQuerySplitResult")
         .allowlist_type("PgQuerySplitStmt")
+        .allowlist_type("PgQueryPlpgsqlProtobufParseResult")
         // Also generate bindings for size_t since it's used in PgQueryProtobuf
         .allowlist_type("size_t")
         .allowlist_var("PG_VERSION_NUM");
@@ -269,6 +272,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
             bindings_content.push_str("    pub fn pg_query_scan(input: *const ::std::os::raw::c_char) -> PgQueryScanResult;\n");
             bindings_content.push_str("    pub fn pg_query_parse_protobuf(input: *const ::std::os::raw::c_char) -> PgQueryProtobufParseResult;\n");
             bindings_content.push_str("    pub fn pg_query_parse_plpgsql(input: *const ::std::os::raw::c_char) -> PgQueryPlpgsqlParseResult;\n");
+            bindings_content.push_str("    pub fn pg_query_parse_plpgsql_protobuf(input: *const ::std::os::raw::c_char) -> PgQueryPlpgsqlProtobufParseResult;\n");
             bindings_content.push_str("    pub fn pg_query_deparse_protobuf(protobuf: PgQueryProtobuf) -> PgQueryDeparseResult;\n");
             bindings_content.push_str("    pub fn pg_query_normalize(input: *const ::std::os::raw::c_char) -> PgQueryNormalizeResult;\n");
             bindings_content.push_str("    pub fn pg_query_fingerprint(input: *const ::std::os::raw::c_char) -> PgQueryFingerprintResult;\n");
@@ -278,6 +282,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 .push_str("    pub fn pg_query_free_scan_result(result: PgQueryScanResult);\n");
             bindings_content.push_str("    pub fn pg_query_free_protobuf_parse_result(result: PgQueryProtobufParseResult);\n");
             bindings_content.push_str("    pub fn pg_query_free_plpgsql_parse_result(result: PgQueryPlpgsqlParseResult);\n");
+            bindings_content.push_str("    pub fn pg_query_free_plpgsql_protobuf_parse_result(result: PgQueryPlpgsqlProtobufParseResult);\n");
             bindings_content.push_str(
                 "    pub fn pg_query_free_deparse_result(result: PgQueryDeparseResult);\n",
             );

crates/pgls_query/src/lib.rs

@@ -28,6 +28,9 @@ pub use scan::*;
 pub use split::*;
 
 pub use protobuf::Node;
+pub use protobuf::PLpgSqlFunction;
+pub use protobuf::PLpgSqlParseResult;
+pub use protobuf::PLpgSqlStmt;
 
 // Include the generated bindings with 2024 edition compatibility
 #[allow(non_upper_case_globals)]

crates/pgls_query/src/plpgsql.rs

@@ -1,15 +1,19 @@
 use std::ffi::{CStr, CString};
 
+use prost::Message;
+
 use crate::bindings::*;
 use crate::error::*;
+use crate::protobuf;
 
-/// An experimental API which parses a PLPGSQL function. This currently drops the returned
-/// structure and returns only a Result<()>.
+/// Parses the given PL/pgSQL function into an abstract syntax tree.
 ///
 /// # Example
 ///
 /// ```rust
-/// let result = pgls_query::parse_plpgsql("
+/// use pgls_query::parse_plpgsql;
+///
+/// let result = parse_plpgsql("
 ///     CREATE OR REPLACE FUNCTION cs_fmt_browser_version(v_name varchar, v_version varchar)
 ///     RETURNS varchar AS $$
 ///     BEGIN
@@ -21,18 +25,255 @@ use crate::error::*;
 ///     $$ LANGUAGE plpgsql;
 /// ");
 /// assert!(result.is_ok());
+/// let result = result.unwrap();
+/// assert_eq!(result.functions().len(), 1);
 /// ```
-pub fn parse_plpgsql(stmt: &str) -> Result<()> {
+pub fn parse_plpgsql(stmt: &str) -> Result<PlpgsqlParseResult> {
     let input = CString::new(stmt)?;
-    let result = unsafe { pg_query_parse_plpgsql(input.as_ptr()) };
-    let structure = if !result.error.is_null() {
+    let result = unsafe { pg_query_parse_plpgsql_protobuf(input.as_ptr()) };
+    let parse_result = if !result.error.is_null() {
         let message = unsafe { CStr::from_ptr((*result.error).message) }
             .to_string_lossy()
             .to_string();
         Err(Error::Parse(message))
     } else {
-        Ok(())
+        let data = unsafe {
+            std::slice::from_raw_parts(
+                result.parse_tree.data as *const u8,
+                result.parse_tree.len as usize,
+            )
+        };
+        protobuf::PLpgSqlParseResult::decode(data)
+            .map_err(Error::Decode)
+            .map(PlpgsqlParseResult::new)
     };
-    unsafe { pg_query_free_plpgsql_parse_result(result) };
-    structure
+    unsafe { pg_query_free_plpgsql_protobuf_parse_result(result) };
+    parse_result
+}
+
+/// The result of parsing a PL/pgSQL function
+#[derive(Debug)]
+pub struct PlpgsqlParseResult {
+    /// The parsed protobuf result
+    pub protobuf: protobuf::PLpgSqlParseResult,
+}
+
+impl PlpgsqlParseResult {
+    /// Create a new PlpgsqlParseResult
+    pub fn new(protobuf: protobuf::PLpgSqlParseResult) -> Self {
+        Self { protobuf }
+    }
+
+    /// Returns a reference to the single parsed PL/pgSQL function.
+    ///
+    /// Returns `None` if there is not exactly one function in the parse result.
+    /// Use `functions()` if you need to handle multiple functions.
+    pub fn function(&self) -> Option<&protobuf::PLpgSqlFunction> {
+        if self.protobuf.plpgsql_funcs.len() != 1 {
+            return None;
+        }
+        self.protobuf.plpgsql_funcs.first()
+    }
+
+    /// Consumes the result and returns the single parsed function.
+    ///
+    /// Returns `None` if there is not exactly one function in the parse result.
+    /// Use `into_functions()` if you need to handle multiple functions.
+    pub fn into_function(self) -> Option<protobuf::PLpgSqlFunction> {
+        if self.protobuf.plpgsql_funcs.len() != 1 {
+            return None;
+        }
+        self.protobuf.plpgsql_funcs.into_iter().next()
+    }
+
+    /// Returns a reference to the list of parsed PL/pgSQL functions
+    pub fn functions(&self) -> &[protobuf::PLpgSqlFunction] {
+        &self.protobuf.plpgsql_funcs
+    }
+
+    /// Consumes the result and returns the list of parsed functions
+    pub fn into_functions(self) -> Vec<protobuf::PLpgSqlFunction> {
+        self.protobuf.plpgsql_funcs
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::protobuf::p_lpg_sql_stmt::Stmt;
+
+    #[test]
+    fn test_parse_plpgsql_simple() {
+        let result = parse_plpgsql(
+            "
+            CREATE OR REPLACE FUNCTION test_func()
+            RETURNS void AS $$
+            BEGIN
+                NULL;
+            END;
+            $$ LANGUAGE plpgsql;
+        ",
+        );
+        assert!(result.is_ok());
+        let result = result.unwrap();
+
+        // Use function() for single function access
+        let func = result.function().expect("should have exactly one function");
+        assert!(func.action.is_some());
+
+        // The body should contain statements
+        let action = func.action.as_ref().unwrap();
+        assert!(!action.body.is_empty());
+    }
+
+    #[test]
+    fn test_parse_plpgsql_with_assignment() {
+        let result = parse_plpgsql(
+            "
+            CREATE OR REPLACE FUNCTION add_numbers(a int, b int)
+            RETURNS int AS $$
+            DECLARE
+                result int;
+            BEGIN
+                result := a + b;
+                RETURN result;
+            END;
+            $$ LANGUAGE plpgsql;
+        ",
+        );
+        assert!(result.is_ok());
+        let result = result.unwrap();
+        assert_eq!(result.functions().len(), 1);
+
+        let func = &result.functions()[0];
+        let action = func.action.as_ref().unwrap();
+
+        // Should have assignment and return statements
+        assert!(action.body.len() >= 2);
+
+        // First statement should be an assignment
+        let first_stmt = &action.body[0];
+        assert!(matches!(first_stmt.stmt, Some(Stmt::StmtAssign(_))));
+
+        // Second statement should be a return
+        let second_stmt = &action.body[1];
+        assert!(matches!(second_stmt.stmt, Some(Stmt::StmtReturn(_))));
+
+        // Verify the assignment expression contains the query
+        if let Some(Stmt::StmtAssign(assign)) = &first_stmt.stmt {
+            assert!(assign.expr.is_some());
+            let expr = assign.expr.as_ref().unwrap();
+            assert!(expr.query.contains("a + b"));
+        }
+    }
+
+    #[test]
+    fn test_parse_plpgsql_with_if() {
+        let result = parse_plpgsql(
+            "
+            CREATE OR REPLACE FUNCTION cs_fmt_browser_version(v_name varchar, v_version varchar)
+            RETURNS varchar AS $$
+            BEGIN
+                IF v_version IS NULL THEN
+                    RETURN v_name;
+                END IF;
+                RETURN v_name || '/' || v_version;
+            END;
+            $$ LANGUAGE plpgsql;
+        ",
+        );
+        assert!(result.is_ok());
+        let result = result.unwrap();
+        assert_eq!(result.functions().len(), 1);
+
+        let func = &result.functions()[0];
+        let action = func.action.as_ref().unwrap();
+
+        // Should have IF and RETURN statements
+        assert!(action.body.len() >= 2);
+
+        // First statement should be IF
+        let if_stmt = &action.body[0];
+        assert!(matches!(if_stmt.stmt, Some(Stmt::StmtIf(_))));
+
+        // Verify the IF statement structure
+        if let Some(Stmt::StmtIf(if_node)) = &if_stmt.stmt {
+            // Should have a condition
+            assert!(if_node.cond.is_some());
+            let cond = if_node.cond.as_ref().unwrap();
+            assert!(cond.query.contains("v_version IS NULL"));
+
+            // Should have a then_body with RETURN statement
+            assert!(!if_node.then_body.is_empty());
+            assert!(matches!(
+                if_node.then_body[0].stmt,
+                Some(Stmt::StmtReturn(_))
+            ));
+        }
+
+        // Second statement should be RETURN
+        let return_stmt = &action.body[1];
+        assert!(matches!(return_stmt.stmt, Some(Stmt::StmtReturn(_))));
+    }
+
+    #[test]
+    fn test_parse_plpgsql_with_loop() {
+        let result = parse_plpgsql(
+            "
+            CREATE OR REPLACE FUNCTION count_down(n int)
+            RETURNS void AS $$
+            BEGIN
+                WHILE n > 0 LOOP
+                    n := n - 1;
+                END LOOP;
+            END;
+            $$ LANGUAGE plpgsql;
+        ",
+        );
+        assert!(result.is_ok());
+        let result = result.unwrap();
+
+        let func = &result.functions()[0];
+        let action = func.action.as_ref().unwrap();
+
+        // First statement should be WHILE loop
+        let while_stmt = &action.body[0];
+        assert!(matches!(while_stmt.stmt, Some(Stmt::StmtWhile(_))));
+
+        if let Some(Stmt::StmtWhile(while_node)) = &while_stmt.stmt {
+            // Should have a condition
+            assert!(while_node.cond.is_some());
+            let cond = while_node.cond.as_ref().unwrap();
+            assert!(cond.query.contains("n > 0"));
+
+            // Should have a body with assignment
+            assert!(!while_node.body.is_empty());
+        }
+    }
+
+    #[test]
+    fn test_parse_plpgsql_error() {
+        let result = parse_plpgsql("not valid plpgsql");
+        assert!(result.is_err());
+    }
+
+    #[test]
+    fn test_parse_plpgsql_multiple_functions() {
+        let result = parse_plpgsql(
+            "
+            CREATE FUNCTION foo() RETURNS void AS $$ BEGIN NULL; END; $$ LANGUAGE plpgsql;
+            CREATE FUNCTION bar() RETURNS int AS $$ BEGIN RETURN 1; END; $$ LANGUAGE plpgsql;
+        ",
+        );
+        assert!(result.is_ok());
+        let result = result.unwrap();
+
+        // function() returns None when multiple functions present
+        assert!(result.function().is_none());
+
+        // Use functions() for multiple
+        assert_eq!(result.functions().len(), 2);
+        assert_eq!(result.functions()[0].fn_signature, "foo");
+        assert_eq!(result.functions()[1].fn_signature, "bar");
+    }
 }