pgdog_plugin/
context.rs

1//! Context passed to and from the plugins.
2
3use std::ops::Deref;
4
5use crate::{bindings::PdRouterContext, PdRoute, PdStatement};
6
7/// PostgreSQL statement, parsed by [`pg_query`].
8///
9/// Implements [`Deref`] on [`PdStatement`], which is passed
10/// in using the FFI interface.
11/// Use the [`PdStatement::protobuf`] method to obtain a reference
12/// to the Abstract Syntax Tree.
13///
14/// ### Example
15///
16/// ```ignore
17/// let ast = statement.protobuf();
18/// println!("{:#?}", ast);
19/// ```
20pub struct Statement {
21    ffi: PdStatement,
22}
23
24impl Deref for Statement {
25    type Target = PdStatement;
26
27    fn deref(&self) -> &Self::Target {
28        &self.ffi
29    }
30}
31
32/// Context information provided by PgDog to the plugin at statement execution. It contains the actual statement and several metadata about
33/// the state of the database cluster:
34///
35/// - Number of shards
36/// - Does it have replicas
37/// - Does it have a primary
38///
39/// ### Example
40///
41/// ```
42/// use pgdog_plugin::{Context, Route, macros, Shard, ReadWrite};
43///
44/// #[macros::route]
45/// fn route(context: Context) -> Route {
46///     let shards = context.shards();
47///     let read_only = context.read_only();
48///     let ast = context.statement().protobuf();
49///
50///     println!("shards: {} (read_only: {})", shards, read_only);
51///     println!("ast: {:#?}", ast);
52///
53///     let read_write = if read_only {
54///         ReadWrite::Read
55///     } else {
56///         ReadWrite::Write
57///     };
58///
59///     Route::new(Shard::Direct(0), read_write)
60/// }
61/// ```
62///
63pub struct Context {
64    ffi: PdRouterContext,
65}
66
67impl From<PdRouterContext> for Context {
68    fn from(value: PdRouterContext) -> Self {
69        Self { ffi: value }
70    }
71}
72
73impl Context {
74    /// Returns a reference to the Abstract Syntax Tree (AST) created by [`pg_query`].
75    ///
76    /// # Example
77    ///
78    /// ```ignore
79    /// let ast = context.statement().protobuf();
80    /// let nodes = ast.nodes();
81    /// ```
82    pub fn statement(&self) -> Statement {
83        Statement {
84            ffi: self.ffi.query,
85        }
86    }
87
88    /// Returns true if the database cluster doesn't have a primary database and can only serve
89    /// read queries.
90    ///
91    /// # Example
92    ///
93    /// ```
94    /// # use pgdog_plugin::Context;
95    /// # let context = unsafe { Context::doc_test() };
96    /// let read_only = context.read_only();
97    ///
98    /// if read_only {
99    ///     println!("Database cluster doesn't have a primary, only replicas.");
100    /// }
101    /// ```
102    pub fn read_only(&self) -> bool {
103        self.ffi.has_primary == 0
104    }
105
106    /// Returns true if the database cluster has replica databases.
107    ///
108    /// # Example
109    ///
110    /// ```
111    /// # use pgdog_plugin::Context;
112    /// # let context = unsafe { Context::doc_test() };
113    /// let has_replicas = context.has_replicas();
114    ///
115    /// if has_replicas {
116    ///     println!("Database cluster can load balance read queries.")
117    /// }
118    /// ```
119    pub fn has_replicas(&self) -> bool {
120        self.ffi.has_replicas == 1
121    }
122
123    /// Returns true if the database cluster has a primary database and can serve write queries.
124    ///
125    /// # Example
126    ///
127    /// ```
128    /// # use pgdog_plugin::Context;
129    /// # let context = unsafe { Context::doc_test() };
130    /// let has_primary = context.has_primary();
131    ///
132    /// if has_primary {
133    ///     println!("Database cluster can serve write queries.");
134    /// }
135    /// ```
136    pub fn has_primary(&self) -> bool {
137        !self.read_only()
138    }
139
140    /// Returns the number of shards in the database cluster.
141    ///
142    /// # Example
143    ///
144    /// ```
145    /// # use pgdog_plugin::Context;
146    /// # let context = unsafe { Context::doc_test() };
147    /// let shards = context.shards();
148    ///
149    /// if shards > 1 {
150    ///     println!("Plugin should consider which shard to route the query to.");
151    /// }
152    /// ```
153    pub fn shards(&self) -> usize {
154        self.ffi.shards as usize
155    }
156
157    /// Returns true if the database cluster has more than one shard.
158    ///
159    /// # Example
160    ///
161    /// ```
162    /// # use pgdog_plugin::Context;
163    /// # let context = unsafe { Context::doc_test() };
164    /// let sharded = context.sharded();
165    /// let shards = context.shards();
166    ///
167    /// if sharded {
168    ///     assert!(shards > 1);
169    /// } else {
170    ///     assert_eq!(shards, 1);
171    /// }
172    /// ```
173    pub fn sharded(&self) -> bool {
174        self.shards() > 1
175    }
176
177    /// Returns true if PgDog strongly believes the statement should be sent to a primary. This indicates
178    /// that the statement is **not** a `SELECT` (e.g. `UPDATE`, `DELETE`, etc.), or a `SELECT` that is very likely to write data to the database, e.g.:
179    ///
180    /// ```sql
181    /// WITH users AS (
182    ///     INSERT INTO users VALUES (1, '[email protected]') RETURNING *
183    /// )
184    /// SELECT * FROM users;
185    /// ```
186    ///
187    /// # Example
188    ///
189    /// ```
190    /// # use pgdog_plugin::Context;
191    /// # let context = unsafe { Context::doc_test() };
192    /// if context.write_override() {
193    ///     println!("We should really send this query to the primary.");
194    /// }
195    /// ```
196    pub fn write_override(&self) -> bool {
197        self.ffi.write_override == 1
198    }
199}
200
201impl Context {
202    /// Used for doc tests only. **Do not use**.
203    ///
204    /// # Safety
205    ///
206    /// Not safe, don't use. We use it for doc tests only.
207    ///
208    pub unsafe fn doc_test() -> Context {
209        use std::{os::raw::c_void, ptr::null};
210
211        Context {
212            ffi: PdRouterContext {
213                shards: 1,
214                has_replicas: 1,
215                has_primary: 1,
216                in_transaction: 0,
217                write_override: 0,
218                query: PdStatement {
219                    version: 1,
220                    len: 0,
221                    data: null::<c_void>() as *mut c_void,
222                },
223            },
224        }
225    }
226}
227
228/// What shard, if any, the statement should be sent to.
229///
230/// ### Example
231///
232/// ```
233/// use pgdog_plugin::Shard;
234///
235/// // Send query to shard 2.
236/// let direct = Shard::Direct(2);
237///
238/// // Send query to all shards.
239/// let cross_shard = Shard::All;
240///
241/// // Let PgDog handle sharding.
242/// let unknown = Shard::Unknown;
243/// ```
244#[derive(Debug, Copy, Clone, PartialEq, Eq)]
245pub enum Shard {
246    /// Direct-to-shard statement, sent to the specified shard only.
247    Direct(usize),
248    /// Send statement to all shards and let PgDog collect and transform the results.
249    All,
250    /// Not clear which shard it should go to, so let PgDog decide.
251    /// Use this if you don't want to handle sharding inside the plugin.
252    Unknown,
253}
254
255impl From<Shard> for i64 {
256    fn from(value: Shard) -> Self {
257        match value {
258            Shard::Direct(value) => value as i64,
259            Shard::All => -1,
260            Shard::Unknown => -2,
261        }
262    }
263}
264
265impl TryFrom<i64> for Shard {
266    type Error = ();
267    fn try_from(value: i64) -> Result<Self, Self::Error> {
268        Ok(if value == -1 {
269            Shard::All
270        } else if value == -2 {
271            Shard::Unknown
272        } else if value >= 0 {
273            Shard::Direct(value as usize)
274        } else {
275            return Err(());
276        })
277    }
278}
279
280impl TryFrom<u8> for ReadWrite {
281    type Error = ();
282
283    fn try_from(value: u8) -> Result<Self, Self::Error> {
284        Ok(if value == 0 {
285            ReadWrite::Write
286        } else if value == 1 {
287            ReadWrite::Read
288        } else if value == 2 {
289            ReadWrite::Unknown
290        } else {
291            return Err(());
292        })
293    }
294}
295
296/// Indicates if the statement is a read or a write. Read statements are sent to a replica, if one is configured.
297/// Write statements are sent to the primary.
298///
299/// ### Example
300///
301/// ```
302/// use pgdog_plugin::ReadWrite;
303///
304/// // The statement should go to a replica.
305/// let read = ReadWrite::Read;
306///
307/// // The statement should go the primary.
308/// let write = ReadWrite::Write;
309///
310/// // Skip and let PgDog decide.
311/// let unknown = ReadWrite::Unknown;
312/// ```
313#[derive(Debug, Copy, Clone, PartialEq, Eq)]
314pub enum ReadWrite {
315    /// Send the statement to a replica, if any are configured.
316    Read,
317    /// Send the statement to the primary.
318    Write,
319    /// Plugin doesn't know if the statement is a read or write. This let's PgDog decide.
320    /// Use this if you don't want to make this decision in the plugin.
321    Unknown,
322}
323
324impl From<ReadWrite> for u8 {
325    fn from(value: ReadWrite) -> Self {
326        match value {
327            ReadWrite::Write => 0,
328            ReadWrite::Read => 1,
329            ReadWrite::Unknown => 2,
330        }
331    }
332}
333
334impl Default for PdRoute {
335    fn default() -> Self {
336        Route::unknown().ffi
337    }
338}
339
340/// Statement route.
341///
342/// PgDog uses this to decide where a query should be sent to. Read statements are sent to a replica,
343/// while write ones are sent the primary. If the cluster has more than one shard, the statement can be
344/// sent to a specific database, or all of them.
345///
346/// ### Example
347///
348/// ```
349/// use pgdog_plugin::{Shard, ReadWrite, Route};
350///
351/// // This sends the query to the primary database of shard 0.
352/// let route = Route::new(Shard::Direct(0), ReadWrite::Write);
353///
354/// // This sends the query to all shards, routing them to a replica
355/// // of each shard, if any are configured.
356/// let route = Route::new(Shard::All, ReadWrite::Read);
357///
358/// // No routing information is available. PgDog will ignore it
359/// // and make its own decision.
360/// let route = Route::unknown();
361pub struct Route {
362    ffi: PdRoute,
363}
364
365impl Default for Route {
366    fn default() -> Self {
367        Self::unknown()
368    }
369}
370
371impl Deref for Route {
372    type Target = PdRoute;
373
374    fn deref(&self) -> &Self::Target {
375        &self.ffi
376    }
377}
378
379impl From<PdRoute> for Route {
380    fn from(value: PdRoute) -> Self {
381        Self { ffi: value }
382    }
383}
384
385impl From<Route> for PdRoute {
386    fn from(value: Route) -> Self {
387        value.ffi
388    }
389}
390
391impl Route {
392    /// Create new route.
393    ///
394    /// # Arguments
395    ///
396    /// * `shard`: Which shard the statement should be sent to.
397    /// * `read_write`: Does the statement read or write data. Read statements are sent to a replica. Write statements are sent to the primary.
398    ///
399    pub fn new(shard: Shard, read_write: ReadWrite) -> Route {
400        Self {
401            ffi: PdRoute {
402                shard: shard.into(),
403                read_write: read_write.into(),
404            },
405        }
406    }
407
408    /// Create new route with no sharding or read/write information.
409    /// Use this if you don't want your plugin to do query routing.
410    /// Plugins that do something else with queries, e.g., logging, metrics,
411    /// can return this route.
412    pub fn unknown() -> Route {
413        Self {
414            ffi: PdRoute {
415                shard: -2,
416                read_write: 2,
417            },
418        }
419    }
420}
pgdog_plugin/context.rs

pgdog_plugin/
context.rs
