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