Support parsing complex queries with deeply nested ASTs

.github/workflows/main.yml

@@ -29,6 +29,9 @@ jobs:
     runs-on: ${{ matrix.os }}
 
     steps:
+      - name: Install protobuf-compiler
+        run: sudo apt-get install -y protobuf-compiler
+
       - uses: actions/checkout@v4
         with:
           submodules: recursive
@@ -68,6 +71,9 @@ jobs:
     name: Check file formatting and style
     runs-on: ubuntu-latest
     steps:
+      - name: Install protobuf-compiler
+        run: sudo apt-get install -y protobuf-compiler
+
       - uses: actions/checkout@v4
         with:
           submodules: recursive

Cargo.toml

@@ -13,15 +13,15 @@ repository = "https://github.com/pganalyze/pg_query.rs"
 
 [dependencies]
 itertools = "0.10.3"
-prost = "0.10.4"
+prost = { git = "https://github.com/pganalyze/prost", branch = "recursion-limit-macro" }
 serde = { version = "1.0.139", features = ["derive"] }
 serde_json = "1.0.82"
 thiserror = "1.0.31"
 
 [build-dependencies]
 bindgen = "0.66.1"
 clippy = { version = "0.0.302", optional = true }
-prost-build = "0.10.4"
+prost-build = { git = "https://github.com/pganalyze/prost", branch = "recursion-limit-macro" }
 fs_extra = "1.2.0"
 cc = "1.0.83"
 glob = "0.3.1"

README.md

@@ -61,6 +61,16 @@ let result = pg_query::parse(query).unwrap();
 assert_eq!(result.truncate(32).unwrap(), "INSERT INTO x (...) VALUES (...)");
 ```
 
+## Caveats
+
+When parsing very complex queries you may run into a stack overflow. This can be worked around by using a thread with a custom stack size ([stdlib](https://doc.rust-lang.org/std/thread/index.html#stack-size), [tokio](https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html#method.thread_stack_size)), or using the stacker crate to resize the main thread's stack:
+
+```rust
+stacker::grow(20 * 1024 * 1024, || pg_query::parse(query))
+```
+
+However, a sufficiently complex query could still run into a stack overflow after you increase the stack size. With some work it may be possible to add an adapter API to the prost crate in order to dynamically increase the stack size as needed like [serde_stacker](https://crates.io/crates/serde_stacker) does (if anyone wants to take that on).
+
 ## Credits
 
 Thanks to [Paul Mason](https://github.com/paupino) for his work on [pg_parse](https://github.com/paupino/pg_parse) that this crate is based on.

build.rs

@@ -65,7 +65,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .write_to_file(out_dir.join("bindings.rs"))?;
 
     // Generate the protobuf definition
-    prost_build::compile_protos(&[&out_protobuf_path.join(LIBRARY_NAME).with_extension("proto")], &[&out_protobuf_path])?;
+    let mut config = prost_build::Config::new();
+    config.recursion_limit("ParseResult", 1000);
+    config.compile_protos(&[&out_protobuf_path.join(LIBRARY_NAME).with_extension("proto")], &[&out_protobuf_path])?;
 
     Ok(())
 }

src/node_enum.rs

@@ -231,6 +231,8 @@ impl NodeEnum {
                     }
                     s.index_params.iter().for_each(|n| {
                         if let Some(NodeEnum::IndexElem(n)) = n.node.as_ref() {
+                            iter.push((n.to_ref(), depth, Context::DDL, false));
+
                             if let Some(n) = n.expr.as_ref().and_then(|n| n.node.as_ref()) {
                                 iter.push((n.to_ref(), depth, Context::DDL, false));
                             }

src/query.rs

@@ -37,7 +37,10 @@ pub fn parse(statement: &str) -> Result<ParseResult> {
     } else {
         let data = unsafe { std::slice::from_raw_parts(result.parse_tree.data as *const u8, result.parse_tree.len as usize) };
         let stderr = unsafe { CStr::from_ptr(result.stderr_buffer) }.to_string_lossy().to_string();
-        protobuf::ParseResult::decode(data).map_err(Error::Decode).map(|result| ParseResult::new(result, stderr))
+        match protobuf::ParseResult::decode(data) {
+            Ok(result) => Ok(ParseResult::new(result, stderr)),
+            Err(error) => Err(Error::Decode(error)),
+        }
     };
     unsafe { pg_query_free_protobuf_parse_result(result) };
     parse_result

tests/parse_tests.rs

@@ -4,6 +4,9 @@
 #[cfg(test)]
 use itertools::sorted;
 
+#[cfg(test)]
+use std::thread::Builder;
+
 use pg_query::{
     parse,
     protobuf::{self, a_const::Val},
@@ -31,15 +34,15 @@ fn it_handles_errors() {
 }
 
 #[test]
-fn it_handles_recursion_error() {
+fn it_handles_recursion_without_error_1() {
     let query = "SELECT a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(a(b))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))))";
-    parse(query).err().unwrap();
-    // TODO: unsure how to unwrap the private fields on a protobuf decode error
-    // assert_eq!(error, Error::Decode("recursion limit reached".into()));
+    let result = Builder::new().stack_size(20 * 1024 * 1024).spawn(move || parse(query)).unwrap().join().unwrap().unwrap();
+    assert_eq!(result.tables().len(), 0);
+    assert_eq!(result.statement_types(), ["SelectStmt"]);
 }
 
 #[test]
-fn it_handles_recursion_without_error() {
+fn it_handles_recursion_without_error_2() {
     // The Ruby version of pg_query fails here because of Ruby protobuf limitations
     let query = r#"SELECT * FROM "t0"
         JOIN "t1" ON (1) JOIN "t2" ON (1) JOIN "t3" ON (1) JOIN "t4" ON (1) JOIN "t5" ON (1)