Skip to main content

pg_query/
query.rs

1use std::ffi::{CStr, CString};
2use std::os::raw::c_char;
3
4use prost::Message;
5
6use crate::bindings::*;
7use crate::error::*;
8use crate::parse_result::ParseResult;
9use crate::protobuf;
10
11/// Represents the resulting fingerprint containing both the raw integer form as well as the
12/// corresponding 16 character hex value.
13pub struct Fingerprint {
14    pub value: u64,
15    pub hex: String,
16}
17
18/// PostgreSQL raw-parser entry mode.
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20#[repr(i32)]
21pub enum ParseMode {
22    Default = 0,
23    TypeName = 1,
24    PlPgSqlExpr = 2,
25    PlPgSqlAssign1 = 3,
26    PlPgSqlAssign2 = 4,
27    PlPgSqlAssign3 = 5,
28}
29
30/// Parses the given SQL statement into the given abstract syntax tree.
31///
32/// # Example
33///
34/// ```rust
35/// use pg_query::{Node, NodeEnum, NodeRef};
36///
37/// let result = pg_query::parse("SELECT * FROM contacts");
38/// assert!(result.is_ok());
39/// let result = result.unwrap();
40/// assert_eq!(result.tables(), vec!["contacts"]);
41/// assert!(matches!(result.protobuf.nodes()[0].0, NodeRef::SelectStmt(_)));
42/// ```
43pub fn parse(statement: &str) -> Result<ParseResult> {
44    parse_with_mode(statement, ParseMode::Default)
45}
46
47/// Parses input using one of PostgreSQL's raw-parser modes.
48///
49/// PL/pgSQL consumers should use the mode recorded on `PLpgSQL_expr`
50/// instead of rewriting an expression or assignment into another SQL form.
51pub fn parse_with_mode(statement: &str, mode: ParseMode) -> Result<ParseResult> {
52    let input = CString::new(statement)?;
53    let result = unsafe { pg_query_parse_protobuf_opts(input.as_ptr(), mode as i32) };
54    let parse_result = if !result.error.is_null() {
55        let message = unsafe { CStr::from_ptr((*result.error).message) }
56            .to_string_lossy()
57            .to_string();
58        Err(Error::Parse(message))
59    } else {
60        let data = unsafe {
61            std::slice::from_raw_parts(
62                result.parse_tree.data as *const u8,
63                result.parse_tree.len as usize,
64            )
65        };
66        let stderr = unsafe { CStr::from_ptr(result.stderr_buffer) }
67            .to_string_lossy()
68            .to_string();
69        protobuf::ParseResult::decode(data)
70            .map_err(Error::Decode)
71            .map(|result| ParseResult::new(result, stderr))
72    };
73    unsafe { pg_query_free_protobuf_parse_result(result) };
74    parse_result
75}
76
77/// Converts a parsed tree back into a string.
78///
79/// # Example
80///
81/// ```rust
82/// use pg_query::{Node, NodeEnum, NodeRef};
83///
84/// let result = pg_query::parse("INSERT INTO other (name) SELECT name FROM contacts");
85/// let result = result.unwrap();
86/// let insert = result.protobuf.nodes()[0].0;
87/// let select = result.protobuf.nodes()[1].0;
88/// assert!(matches!(insert, NodeRef::InsertStmt(_)));
89/// assert!(matches!(select, NodeRef::SelectStmt(_)));
90///
91/// // The entire parse result can be deparsed:
92/// assert_eq!(result.deparse().unwrap(), "INSERT INTO other (name) SELECT name FROM contacts");
93/// // Or an individual node can be deparsed:
94/// assert_eq!(insert.deparse().unwrap(), "INSERT INTO other (name) SELECT name FROM contacts");
95/// assert_eq!(select.deparse().unwrap(), "SELECT name FROM contacts");
96/// ```
97///
98/// Note that this function will panic if called on a node not defined in `deparseStmt`
99pub fn deparse(protobuf: &protobuf::ParseResult) -> Result<String> {
100    let buffer = protobuf.encode_to_vec();
101    let len = buffer.len();
102    let data = buffer.as_ptr() as *const c_char as *mut c_char;
103    let protobuf = PgQueryProtobuf { data, len };
104    let result = unsafe { pg_query_deparse_protobuf(protobuf) };
105
106    let deparse_result = if !result.error.is_null() {
107        let message = unsafe { CStr::from_ptr((*result.error).message) }
108            .to_string_lossy()
109            .to_string();
110        Err(Error::Parse(message))
111    } else {
112        let query = unsafe { CStr::from_ptr(result.query) }
113            .to_string_lossy()
114            .to_string();
115        Ok(query)
116    };
117
118    unsafe { pg_query_free_deparse_result(result) };
119    deparse_result
120}
121
122/// Normalizes the given SQL statement, returning a parametized version.
123///
124/// # Example
125///
126/// ```rust
127/// let result = pg_query::normalize("SELECT * FROM contacts WHERE name='Paul'");
128/// assert!(result.is_ok());
129/// let result = result.unwrap();
130/// assert_eq!(result, "SELECT * FROM contacts WHERE name=$1");
131/// ```
132pub fn normalize(statement: &str) -> Result<String> {
133    let input = CString::new(statement)?;
134    let result = unsafe { pg_query_normalize(input.as_ptr()) };
135    let normalized_query = if !result.error.is_null() {
136        let message = unsafe { CStr::from_ptr((*result.error).message) }
137            .to_string_lossy()
138            .to_string();
139        Err(Error::Parse(message))
140    } else {
141        let n = unsafe { CStr::from_ptr(result.normalized_query) };
142        Ok(n.to_string_lossy().to_string())
143    };
144    unsafe { pg_query_free_normalize_result(result) };
145    normalized_query
146}
147
148/// Fingerprints the given SQL statement. Useful for comparing parse trees across different implementations
149/// of `libpg_query`.
150///
151/// # Example
152///
153/// ```rust
154/// let result = pg_query::fingerprint("SELECT * FROM contacts WHERE name='Paul'");
155/// assert!(result.is_ok());
156/// let result = result.unwrap();
157/// assert_eq!(result.hex, "0e2581a461ece536");
158/// ```
159pub fn fingerprint(statement: &str) -> Result<Fingerprint> {
160    let input = CString::new(statement)?;
161    let result = unsafe { pg_query_fingerprint(input.as_ptr()) };
162    let fingerprint = if !result.error.is_null() {
163        let message = unsafe { CStr::from_ptr((*result.error).message) }
164            .to_string_lossy()
165            .to_string();
166        Err(Error::Parse(message))
167    } else {
168        let hex = unsafe { CStr::from_ptr(result.fingerprint_str) };
169        Ok(Fingerprint {
170            value: result.fingerprint,
171            hex: hex.to_string_lossy().to_string(),
172        })
173    };
174    unsafe { pg_query_free_fingerprint_result(result) };
175    fingerprint
176}
177
178/// An experimental API which parses a PLPGSQL function. This currently returns the raw JSON structure.
179///
180/// # Example
181///
182/// ```rust
183/// let result = pg_query::parse_plpgsql("
184///     CREATE OR REPLACE FUNCTION cs_fmt_browser_version(v_name varchar, v_version varchar)
185///     RETURNS varchar AS $$
186///     BEGIN
187///         IF v_version IS NULL THEN
188///             RETURN v_name;
189///         END IF;
190///         RETURN v_name || '/' || v_version;
191///     END;
192///     $$ LANGUAGE plpgsql;
193/// ");
194/// assert!(result.is_ok());
195/// ```
196pub fn parse_plpgsql(stmt: &str) -> Result<serde_json::Value> {
197    let input = CString::new(stmt)?;
198    let result = unsafe { pg_query_parse_plpgsql(input.as_ptr()) };
199    let structure = if !result.error.is_null() {
200        let message = unsafe { CStr::from_ptr((*result.error).message) }
201            .to_string_lossy()
202            .to_string();
203        Err(Error::Parse(message))
204    } else {
205        let raw = unsafe { CStr::from_ptr(result.plpgsql_funcs) };
206        serde_json::from_str(&raw.to_string_lossy()).map_err(|e| Error::InvalidJson(e.to_string()))
207    };
208    unsafe { pg_query_free_plpgsql_parse_result(result) };
209    structure
210}
211
212/// Split a well-formed query into separate statements.
213///
214/// # Example
215///
216/// ```rust
217/// let query = r#"select /*;*/ 1; select "2;", (select 3);"#;
218/// let statements = pg_query::split_with_parser(query).unwrap();
219/// assert_eq!(statements, vec!["select /*;*/ 1", r#"select "2;", (select 3)"#]);
220/// ```
221///
222/// However, `split_with_parser` will fail on malformed statements
223///
224/// ```rust
225/// let query = "select 1; this statement is not sql; select 2;";
226/// let result = pg_query::split_with_parser(query);
227/// let err = r#"syntax error at or near "this""#;
228/// assert_eq!(result, Err(pg_query::Error::Split(err.to_string())));
229/// ```
230pub fn split_with_parser(query: &str) -> Result<Vec<&str>> {
231    let input = CString::new(query)?;
232    let result = unsafe { pg_query_split_with_parser(input.as_ptr()) };
233    let split_result = if !result.error.is_null() {
234        let message = unsafe { CStr::from_ptr((*result.error).message) }
235            .to_string_lossy()
236            .to_string();
237        Err(Error::Split(message))
238    } else {
239        let n_stmts = result.n_stmts as usize;
240        let mut statements = Vec::with_capacity(n_stmts);
241        for offset in 0..n_stmts {
242            let split_stmt = unsafe { *result.stmts.add(offset).read() };
243            let start = split_stmt.stmt_location as usize;
244            let end = start + split_stmt.stmt_len as usize;
245            statements.push(&query[start..end]);
246            // not sure the start..end slice'll hold up for non-utf8 charsets
247        }
248        Ok(statements)
249    };
250    unsafe { pg_query_free_split_result(result) };
251    split_result
252}
253
254/// Scan a sql query into a its component of tokens.
255///
256/// # Example
257///
258/// ```rust
259/// use pg_query::protobuf::*;
260/// let sql = "SELECT update AS left /* comment */ FROM between";
261/// let result = pg_query::scan(sql).unwrap();
262/// let tokens: Vec<std::string::String> = result.tokens.iter().map(|token| {
263///     format!("{:?}", token)
264/// }).collect();
265/// assert_eq!(
266///     tokens,
267///     vec![
268///         "ScanToken { start: 0, end: 6, token: Select, keyword_kind: ReservedKeyword }",
269///         "ScanToken { start: 7, end: 13, token: Update, keyword_kind: UnreservedKeyword }",
270///         "ScanToken { start: 14, end: 16, token: As, keyword_kind: ReservedKeyword }",
271///         "ScanToken { start: 17, end: 21, token: Left, keyword_kind: TypeFuncNameKeyword }",
272///         "ScanToken { start: 22, end: 35, token: CComment, keyword_kind: NoKeyword }",
273///         "ScanToken { start: 36, end: 40, token: From, keyword_kind: ReservedKeyword }",
274///         "ScanToken { start: 41, end: 48, token: Between, keyword_kind: ColNameKeyword }"
275///     ]);
276/// ```
277pub fn scan(sql: &str) -> Result<protobuf::ScanResult> {
278    let input = CString::new(sql)?;
279    let result = unsafe { pg_query_scan(input.as_ptr()) };
280    let scan_result = if !result.error.is_null() {
281        let message = unsafe { CStr::from_ptr((*result.error).message) }
282            .to_string_lossy()
283            .to_string();
284        Err(Error::Scan(message))
285    } else {
286        let data = unsafe {
287            std::slice::from_raw_parts(result.pbuf.data as *const u8, result.pbuf.len as usize)
288        };
289        protobuf::ScanResult::decode(data).map_err(Error::Decode)
290    };
291    unsafe { pg_query_free_scan_result(result) };
292    scan_result
293}
294
295/// Split a potentially-malformed query into separate statements. Note that
296/// invalid tokens will be skipped
297/// ```rust
298/// let query = r#"select /*;*/ 1; asdf; select "2;", (select 3); asdf"#;
299/// let statements = pg_query::split_with_scanner(query).unwrap();
300/// assert_eq!(statements, vec![
301///     "select /*;*/ 1",
302///     // skipped " asdf" since it was an invalid token
303///     r#" select "2;", (select 3)"#,
304/// ]);
305/// ```
306pub fn split_with_scanner(query: &str) -> Result<Vec<&str>> {
307    let input = CString::new(query)?;
308    let result = unsafe { pg_query_split_with_scanner(input.as_ptr()) };
309    let split_result = if !result.error.is_null() {
310        let message = unsafe { CStr::from_ptr((*result.error).message) }
311            .to_string_lossy()
312            .to_string();
313        Err(Error::Split(message))
314    } else {
315        // don't use result.stderr_buffer since it appears unused unless
316        // libpg_query is compiled with DEBUG defined.
317        let n_stmts = result.n_stmts as usize;
318        let mut start: usize;
319        let mut end: usize;
320        let mut statements = Vec::with_capacity(n_stmts);
321        for offset in 0..n_stmts {
322            let split_stmt = unsafe { *result.stmts.add(offset).read() };
323            start = split_stmt.stmt_location as usize;
324            // TODO: consider comparing the new value of start to the old value
325            // of end to see if any region larger than a statement-separator got skipped
326            end = start + split_stmt.stmt_len as usize;
327            statements.push(&query[start..end]);
328        }
329        Ok(statements)
330    };
331    unsafe { pg_query_free_split_result(result) };
332    split_result
333}