pgdog_plugin/
ast.rs

1//! Wrapper around `pg_query` protobuf-generated statement.
2//!
3//! This is passed through FFI safely by ensuring two conditions:
4//!
5//! 1. The version of the **Rust compiler** used to build the plugin is the same used to build PgDog
6//! 2. The version of the **`pg_query` library** used by the plugin is the same used by PgDog
7//!
8use std::{ffi::c_void, ops::Deref};
9
10use pg_query::protobuf::{ParseResult, RawStmt};
11
12use crate::bindings::PdStatement;
13
14impl PdStatement {
15    /// Create FFI binding from `pg_query` output.
16    ///
17    /// # Safety
18    ///
19    /// The reference must live for the entire time
20    /// this struct is used. This is _not_ checked by the compiler,
21    /// and is the responsibility of the caller.
22    ///
23    pub unsafe fn from_proto(value: &ParseResult) -> Self {
24        Self {
25            data: value.stmts.as_ptr() as *mut c_void,
26            version: value.version,
27            len: value.stmts.len() as u64,
28        }
29    }
30}
31
32/// Wrapper around [`pg_query::protobuf::ParseResult`], which allows
33/// the caller to use `pg_query` types and methods to inspect the statement.
34#[derive(Debug)]
35pub struct PdParseResult {
36    parse_result: Option<ParseResult>,
37    borrowed: bool,
38}
39
40impl Clone for PdParseResult {
41    /// Cloning the binding is safe. A new structure
42    /// will be created without any references to the original
43    /// reference.
44    fn clone(&self) -> Self {
45        Self {
46            parse_result: self.parse_result.clone(),
47            borrowed: false,
48        }
49    }
50}
51
52impl From<PdStatement> for PdParseResult {
53    /// Create the binding from a FFI-passed reference.
54    ///
55    /// SAFETY: Memory is owned by the caller.
56    ///
57    fn from(value: PdStatement) -> Self {
58        Self {
59            parse_result: Some(ParseResult {
60                version: value.version,
61                stmts: unsafe {
62                    Vec::from_raw_parts(
63                        value.data as *mut RawStmt,
64                        value.len as usize,
65                        value.len as usize,
66                    )
67                },
68            }),
69            borrowed: true,
70        }
71    }
72}
73
74impl Drop for PdParseResult {
75    /// Drop the binding and forget the memory if the binding
76    /// is using the referenced struct. Otherwise, deallocate as normal.
77    fn drop(&mut self) {
78        if self.borrowed {
79            let parse_result = self.parse_result.take();
80            std::mem::forget(parse_result.unwrap().stmts);
81        }
82    }
83}
84
85impl Deref for PdParseResult {
86    type Target = ParseResult;
87
88    fn deref(&self) -> &Self::Target {
89        self.parse_result.as_ref().unwrap()
90    }
91}
92
93impl PdStatement {
94    /// Get the protobuf-wrapped PostgreSQL statement. The returned structure, [`PdParseResult`],
95    /// implements [`Deref`] to [`pg_query::protobuf::ParseResult`] and can be used to parse the
96    /// statement.
97    pub fn protobuf(&self) -> PdParseResult {
98        PdParseResult::from(*self)
99    }
100}
101
102#[cfg(test)]
103mod test {
104    use pg_query::NodeEnum;
105
106    use super::*;
107
108    #[test]
109    fn test_ast() {
110        let ast = pg_query::parse("SELECT * FROM users WHERE id = $1").unwrap();
111        let ffi = unsafe { PdStatement::from_proto(&ast.protobuf) };
112        match ffi
113            .protobuf()
114            .stmts
115            .first()
116            .unwrap()
117            .stmt
118            .as_ref()
119            .unwrap()
120            .node
121            .as_ref()
122            .unwrap()
123        {
124            NodeEnum::SelectStmt(_) => (),
125            _ => {
126                panic!("not a select")
127            }
128        };
129
130        let _ = ffi.protobuf().clone();
131    }
132}
pgdog_plugin/ast.rs

pgdog_plugin/
ast.rs
