Add pg_query::summary() (#62)

Author: duckinator

Cargo.lock

@@ -34,6 +34,11 @@ cc = "1.0.83"
 glob = "0.3.1"
 
 [dev-dependencies]
+brunch = "0.11"
 easy-parallel = "3.2.0"
 pretty_assertions = "1.4.0"
 regex = "1.6.0"
+
+[[bench]]
+name = "parse_vs_summary"
+harness = false

Cargo.toml

@@ -0,0 +1,26 @@
+use brunch::Bench;
+use pg_query;
+
+brunch::benches!(
+    Bench::new("parse").run_seeded_with(seed, |query| pg_query::parse(&query)),
+    Bench::new("summary").run_seeded_with(seed, |query| pg_query::summary(&query, -1)),
+    // I had to be less mean just so the parse+truncate one didn't crash.
+    Bench::new("parse + truncate").run_seeded_with(less_mean_seed, |query| pg_query::parse(&query).unwrap().truncate(50).unwrap()),
+    Bench::new("summary + truncate").run_seeded_with(less_mean_seed, |query| pg_query::summary(&query, 50)),
+);
+
+fn less_mean_seed() -> String {
+    build_query(30)
+}
+
+fn seed() -> String {
+    build_query(500)
+}
+
+fn build_query(table_references: i32) -> String {
+    let mut query = "SELECT * FROM t".to_string();
+    for i in 0..table_references {
+        query = format!("{query} JOIN t{i} ON t.id = t{i}.t_id AND t{i}.k IN (1, 2, 3, 4) AND t{i}.f IN (SELECT o FROM p WHERE q = 'foo')");
+    }
+    query
+}

benches/parse_vs_summary.rs

@@ -21,6 +21,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Copy the relevant source files to the OUT_DIR
     let source_paths = vec![
         build_path.join("pg_query").with_extension("h"),
+        build_path.join("postgres_deparse").with_extension("h"),
         build_path.join("Makefile"),
         build_path.join("src"),
         build_path.join("protobuf"),

build.rs

@@ -48,6 +48,8 @@ mod parse_result;
 #[rustfmt::skip]
 pub mod protobuf;
 mod query;
+mod summary;
+mod summary_result;
 mod truncate;
 
 pub use error::*;
@@ -56,6 +58,8 @@ pub use node_mut::*;
 pub use node_ref::*;
 pub use parse_result::*;
 pub use query::*;
+pub use summary::*;
+pub use summary_result::*;
 pub use truncate::*;
 
 pub use protobuf::Node;

libpg_query

@@ -4455,6 +4455,121 @@ pub struct ScanToken {
     #[prost(enumeration = "KeywordKind", tag = "5")]
     pub keyword_kind: i32,
 }
+/// protobuf-c doesn't support optional fields, so any optional strings
+/// are just an empty string if it should be the equivalent of None/nil.
+///
+/// These fields have `// optional` at the end of the line.
+///
+/// Upstream issue: <https://github.com/protobuf-c/protobuf-c/issues/476>
+#[derive(serde::Serialize)]
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct SummaryResult {
+    #[prost(message, repeated, tag = "1")]
+    pub tables: ::prost::alloc::vec::Vec<summary_result::Table>,
+    /// The value here is the table name (i.e. schema.table or just table).
+    #[prost(map = "string, string", tag = "2")]
+    pub aliases: ::std::collections::HashMap<
+        ::prost::alloc::string::String,
+        ::prost::alloc::string::String,
+    >,
+    #[prost(string, repeated, tag = "3")]
+    pub cte_names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
+    #[prost(message, repeated, tag = "4")]
+    pub functions: ::prost::alloc::vec::Vec<summary_result::Function>,
+    #[prost(message, repeated, tag = "5")]
+    pub filter_columns: ::prost::alloc::vec::Vec<summary_result::FilterColumn>,
+    #[prost(string, repeated, tag = "6")]
+    pub statement_types: ::prost::alloc::vec::Vec<::prost::alloc::string::String>,
+    /// optional, empty if truncation limit is -1
+    #[prost(string, tag = "7")]
+    pub truncated_query: ::prost::alloc::string::String,
+}
+/// Nested message and enum types in `SummaryResult`.
+pub mod summary_result {
+    #[derive(serde::Serialize)]
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct Table {
+        #[prost(string, tag = "1")]
+        pub name: ::prost::alloc::string::String,
+        #[prost(string, tag = "2")]
+        pub schema_name: ::prost::alloc::string::String,
+        #[prost(string, tag = "3")]
+        pub table_name: ::prost::alloc::string::String,
+        #[prost(enumeration = "Context", tag = "4")]
+        pub context: i32,
+    }
+    #[derive(serde::Serialize)]
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct Function {
+        #[prost(string, tag = "1")]
+        pub name: ::prost::alloc::string::String,
+        #[prost(string, tag = "2")]
+        pub function_name: ::prost::alloc::string::String,
+        /// optional
+        #[prost(string, tag = "3")]
+        pub schema_name: ::prost::alloc::string::String,
+        #[prost(enumeration = "Context", tag = "4")]
+        pub context: i32,
+    }
+    #[derive(serde::Serialize)]
+    #[derive(Clone, PartialEq, ::prost::Message)]
+    pub struct FilterColumn {
+        /// optional
+        #[prost(string, tag = "1")]
+        pub schema_name: ::prost::alloc::string::String,
+        /// optional
+        #[prost(string, tag = "2")]
+        pub table_name: ::prost::alloc::string::String,
+        #[prost(string, tag = "3")]
+        pub column: ::prost::alloc::string::String,
+    }
+    #[derive(serde::Serialize)]
+    #[derive(
+        Clone,
+        Copy,
+        Debug,
+        PartialEq,
+        Eq,
+        Hash,
+        PartialOrd,
+        Ord,
+        ::prost::Enumeration
+    )]
+    #[repr(i32)]
+    pub enum Context {
+        None = 0,
+        Select = 1,
+        Dml = 2,
+        Ddl = 3,
+        Call = 4,
+    }
+    impl Context {
+        /// String value of the enum field names used in the ProtoBuf definition.
+        ///
+        /// The values are not transformed in any way and thus are considered stable
+        /// (if the ProtoBuf definition does not change) and safe for programmatic use.
+        pub fn as_str_name(&self) -> &'static str {
+            match self {
+                Self::None => "None",
+                Self::Select => "Select",
+                Self::Dml => "DML",
+                Self::Ddl => "DDL",
+                Self::Call => "Call",
+            }
+        }
+        /// Creates an enum from field names used in the ProtoBuf definition.
+        pub fn from_str_name(value: &str) -> ::core::option::Option<Self> {
+            match value {
+                "None" => Some(Self::None),
+                "Select" => Some(Self::Select),
+                "DML" => Some(Self::Dml),
+                "DDL" => Some(Self::Ddl),
+                "Call" => Some(Self::Call),
+                _ => None,
+            }
+        }
+    }
+}
 #[derive(serde::Serialize)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)]
 #[repr(i32)]

src/lib.rs

@@ -0,0 +1,46 @@
+use std::ffi::{CStr, CString};
+
+use prost::Message;
+
+use crate::bindings::*;
+use crate::error::*;
+use crate::protobuf;
+use crate::summary_result::SummaryResult;
+
+/// Parses the given SQL statement and provides a summary of it.
+///
+/// It is possible to generate the same data using `pg_query::parse` and
+/// iterating through the parse tree. However, `pg_query::summary` uses a
+/// C implementation to avoid sending as much data over protobuf.
+///
+/// Avoiding sending the parse tree over protobuf can cause as much as an
+/// *order of magnitude* performance improvement. It also prevents some
+/// crashes caused by protobuf handling such a large amount of data.
+///
+/// You can run `cargo bench parse_vs_summary` to run the benchmarks that
+/// comparse the two options.
+///
+/// # Example
+///
+/// ```rust
+/// use pg_query::{Node, NodeEnum, NodeRef};
+///
+/// let result = pg_query::summary("SELECT * FROM contacts", -1);
+/// assert!(result.is_ok());
+/// let result = result.unwrap();
+/// assert_eq!(result.tables(), vec!["contacts"]);
+/// ```
+pub fn summary(statement: &str, truncate_limit: i32) -> Result<SummaryResult> {
+    let input = CString::new(statement)?;
+    let result = unsafe { pg_query_summary(input.as_ptr(), 0, truncate_limit) };
+    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 {
+        let data = unsafe { std::slice::from_raw_parts(result.summary.data as *const u8, result.summary.len as usize) };
+        let stderr = unsafe { CStr::from_ptr(result.stderr_buffer) }.to_string_lossy().to_string();
+        protobuf::SummaryResult::decode(data).map_err(Error::Decode).map(|result| SummaryResult::new(result, stderr))
+    };
+    unsafe { pg_query_free_summary_parse_result(result) };
+    parse_result
+}