pgrx_pg_sys/submodules/pg_try.rs
1//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
2//LICENSE
3//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
4//LICENSE
5//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <contact@pgcentral.org>
6//LICENSE
7//LICENSE All rights reserved.
8//LICENSE
9//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
10use crate::errcodes::PgSqlErrorCode;
11use crate::panic::{downcast_panic_payload, CaughtError};
12use std::collections::BTreeMap;
13use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe, RefUnwindSafe, UnwindSafe};
14
15/// [`PgTryBuilder`] is a mechanism to mimic Postgres C macros `PG_TRY` and `PG_CATCH`.
16///
17/// A primary difference is that the [`PgTryBuilder::finally()`] block runs even if a catch handler
18/// rethrows (or throws a new) error.
19pub struct PgTryBuilder<'a, R, F: FnOnce() -> R + UnwindSafe> {
20 func: F,
21 when: BTreeMap<
22 PgSqlErrorCode,
23 Box<dyn FnMut(CaughtError) -> R + 'a + UnwindSafe + RefUnwindSafe>,
24 >,
25 others: Option<Box<dyn FnMut(CaughtError) -> R + 'a + UnwindSafe + RefUnwindSafe>>,
26 rust: Option<Box<dyn FnMut(CaughtError) -> R + 'a + UnwindSafe + RefUnwindSafe>>,
27 finally: Option<Box<dyn FnMut() + 'a>>,
28}
29
30impl<'a, R, F: FnOnce() -> R + UnwindSafe> PgTryBuilder<'a, R, F> {
31 /// Create a new `[PgTryBuilder]`. The `func` argument specifies the closure that is to run.
32 ///
33 /// If it fails with either a Rust panic or a Postgres error, a registered catch handler
34 /// for that specific error is run. Whether one exists or not, the finally block also runs
35 /// at the end, for any necessary cleanup.
36 ///
37 /// ## Example
38 ///
39 /// ```rust,no_run
40 /// # use pgrx_pg_sys::{ereport, PgTryBuilder};
41 /// # use pgrx_pg_sys::errcodes::PgSqlErrorCode;
42 ///
43 /// let i = 41;
44 /// let mut finished = false;
45 /// let result = PgTryBuilder::new(|| {
46 /// if i < 42 {
47 /// ereport!(ERROR, PgSqlErrorCode::ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE, "number too small");
48 /// }
49 /// i
50 /// })
51 /// .catch_when(PgSqlErrorCode::ERRCODE_NUMERIC_VALUE_OUT_OF_RANGE, |cause| cause.rethrow())
52 /// .finally(|| finished = true)
53 /// .execute();
54 ///
55 /// assert_eq!(finished, true);
56 /// assert_eq!(result, 42);
57 /// ```
58 #[must_use = "must call `PgTryBuilder::execute(self)` in order for it to run"]
59 pub fn new(func: F) -> Self {
60 Self { func, when: Default::default(), others: None, rust: None, finally: None }
61 }
62
63 /// Add a catch handler to run should a specific error occur during execution.
64 ///
65 /// The argument to the catch handler closure is a [`CaughtError`] which can be
66 /// rethrown via [`CaughtError::rethrow()`]
67 ///
68 /// The return value must be of the same type as the main execution block, or the supplied
69 /// error cause must be rethrown.
70 ///
71 /// ## Safety
72 ///
73 /// While this function isn't itself unsafe, catching and then ignoring an important internal
74 /// Postgres error may very well leave your database in an undesirable state. This is your
75 /// responsibility.
76 #[must_use = "must call `PgTryBuilder::execute(self)` in order for it to run"]
77 pub fn catch_when(
78 mut self,
79 error: PgSqlErrorCode,
80 f: impl FnMut(CaughtError) -> R + 'a + UnwindSafe + RefUnwindSafe,
81 ) -> Self {
82 self.when.insert(error, Box::new(f));
83 self
84 }
85
86 /// Add a catch-all handler to catch a raised error that wasn't explicitly caught via
87 /// [PgTryBuilder::catch_when].
88 ///
89 /// The argument to the catch handler closure is a [`CaughtError`] which can be
90 /// rethrown via [`CaughtError::rethrow()`].
91 ///
92 /// The return value must be of the same type as the main execution block, or the supplied
93 /// error cause must be rethrown.
94 ///
95 /// ## Safety
96 ///
97 /// While this function isn't itself unsafe, catching and then ignoring an important internal
98 /// Postgres error may very well leave your database in an undesirable state. This is your
99 /// responsibility.
100 #[must_use = "must call `PgTryBuilder::execute(self)` in order for it to run"]
101 pub fn catch_others(
102 mut self,
103 f: impl FnMut(CaughtError) -> R + 'a + UnwindSafe + RefUnwindSafe,
104 ) -> Self {
105 self.others = Some(Box::new(f));
106 self
107 }
108
109 /// Add a handler to specifically respond to a Rust panic.
110 ///
111 /// The catch handler's closure argument is a [`CaughtError::PostgresError {ereport, payload}`]
112 /// that you can inspect in whatever way makes sense.
113 ///
114 /// The return value must be of the same type as the main execution block, or the supplied
115 /// error cause must be rethrown.
116 ///
117 #[must_use = "must call `PgTryBuilder::execute(self)` in order for it to run"]
118 pub fn catch_rust_panic(
119 mut self,
120 f: impl FnMut(CaughtError) -> R + 'a + UnwindSafe + RefUnwindSafe,
121 ) -> Self {
122 self.rust = Some(Box::new(f));
123 self
124 }
125
126 /// The finally block, of which there can be only one. Successive calls to this function
127 /// will replace the prior finally block.
128 ///
129 /// The finally block closure is called after successful return from the main execution handler
130 /// or a catch handler. The finally block does not return a value.
131 #[must_use = "must call `PgTryBuilder::execute(self)` in order for it to run"]
132 pub fn finally(mut self, f: impl FnMut() + 'a) -> Self {
133 self.finally = Some(Box::new(f));
134 self
135 }
136
137 /// Run the main execution block closure. Any error raised will be passed to a registered
138 /// catch handler, and when finished, the finally block will be run.
139 pub fn execute(mut self) -> R {
140 let result = catch_unwind(self.func);
141
142 fn finally<F: FnMut()>(f: &mut Option<F>) {
143 if let Some(f) = f {
144 f()
145 }
146 }
147
148 let result = match result {
149 Ok(result) => result,
150 Err(error) => {
151 let (sqlerrcode, root_cause) = match downcast_panic_payload(error) {
152 CaughtError::RustPanic { ereport, payload } => {
153 let sqlerrcode = ereport.inner.sqlerrcode;
154 let panic = CaughtError::RustPanic { ereport, payload };
155 (sqlerrcode, panic)
156 }
157 CaughtError::ErrorReport(ereport) => {
158 let sqlerrcode = ereport.inner.sqlerrcode;
159 let panic = CaughtError::ErrorReport(ereport);
160 (sqlerrcode, panic)
161 }
162 CaughtError::PostgresError(ereport) => {
163 let sqlerrcode = ereport.inner.sqlerrcode;
164 let panic = CaughtError::PostgresError(ereport);
165 (sqlerrcode, panic)
166 }
167 };
168
169 // Postgres source docs says that a PG_TRY/PG_CATCH/PG_FINALLY block can't have
170 // both a CATCH and a FINALLY.
171 //
172 // We allow it by wrapping handler execution in its own `catch_unwind()` block
173 // and deferring the finally block execution until after it is complete
174 let handler_result = catch_unwind(AssertUnwindSafe(|| {
175 if let Some(mut handler) = self.when.remove(&sqlerrcode) {
176 // we have a registered catch handler for the error code we caught
177 return handler(root_cause);
178 } else if let Some(mut handler) = self.others {
179 // we have a registered "catch others" handler
180 return handler(root_cause);
181 } else if let Some(mut handler) = self.rust {
182 // we have a registered catch handler for a rust panic
183 if let cause @ CaughtError::RustPanic { .. } = root_cause {
184 // and we have a rust panic
185 return handler(cause);
186 }
187 }
188
189 // we have no handler capable of handling whatever error we have, so rethrow the root cause
190 root_cause.rethrow();
191 }));
192
193 let handler_result = match handler_result {
194 Ok(result) => result,
195 Err(caught) => {
196 let catch_handler_error = downcast_panic_payload(caught);
197
198 // make sure to run the finally block and then resume unwinding
199 // with this new panic
200 finally(&mut self.finally);
201 resume_unwind(Box::new(catch_handler_error))
202 }
203 };
204
205 // Being here means the catch handler didn't raise an error.
206 //
207 // Postgres says:
208 unsafe {
209 /*
210 * FlushErrorState --- flush the error state after error recovery
211 *
212 * This should be called by an error handler after it's done processing
213 * the error; or as soon as it's done CopyErrorData, if it intends to
214 * do stuff that is likely to provoke another error. You are not "out" of
215 * the error subsystem until you have done this.
216 */
217 crate::FlushErrorState();
218 }
219 handler_result
220 }
221 };
222
223 // `result` could be from a successful execution or returned from a catch handler
224 //
225 // Either way, the finally block needs to be run before we return it
226 finally(&mut self.finally);
227 result
228 }
229}