more pg try catching and error handling

This commit is contained in:
2026-03-25 19:31:51 -04:00
parent 4941dc6069
commit ffd6c27da3
5 changed files with 135 additions and 77 deletions

View File

@ -9,6 +9,61 @@ impl SpiExecutor {
pub fn new() -> Self {
Self {}
}
fn transact<F, R>(&self, f: F) -> Result<R, String>
where
F: FnOnce() -> Result<R, String>,
{
unsafe {
let oldcontext = pgrx::pg_sys::CurrentMemoryContext;
let oldowner = pgrx::pg_sys::CurrentResourceOwner;
pgrx::pg_sys::BeginInternalSubTransaction(std::ptr::null());
pgrx::pg_sys::MemoryContextSwitchTo(oldcontext);
let runner = std::panic::AssertUnwindSafe(move || {
let res = f();
pgrx::pg_sys::ReleaseCurrentSubTransaction();
pgrx::pg_sys::MemoryContextSwitchTo(oldcontext);
pgrx::pg_sys::CurrentResourceOwner = oldowner;
res
});
pgrx::PgTryBuilder::new(runner)
.catch_rust_panic(|cause| {
pgrx::pg_sys::RollbackAndReleaseCurrentSubTransaction();
pgrx::pg_sys::MemoryContextSwitchTo(oldcontext);
pgrx::pg_sys::CurrentResourceOwner = oldowner;
// Rust panics are fatal bugs, not validation errors. Rethrow so they bubble up.
cause.rethrow()
})
.catch_others(|cause| {
pgrx::pg_sys::RollbackAndReleaseCurrentSubTransaction();
pgrx::pg_sys::MemoryContextSwitchTo(oldcontext);
pgrx::pg_sys::CurrentResourceOwner = oldowner;
let error_msg = match &cause {
pgrx::pg_sys::panic::CaughtError::PostgresError(e)
| pgrx::pg_sys::panic::CaughtError::ErrorReport(e) => {
let json_err = serde_json::json!({
"error": e.message(),
"code": format!("{:?}", e.sql_error_code()),
"detail": e.detail(),
"hint": e.hint()
});
json_err.to_string()
}
_ => format!("{:?}", cause),
};
pgrx::warning!("JSPG Caught Native Postgres Error: {}", error_msg);
Err(error_msg)
})
.execute()
}
}
}
impl DatabaseExecutor for SpiExecutor {
@ -24,7 +79,7 @@ impl DatabaseExecutor for SpiExecutor {
}
}
pgrx::PgTryBuilder::new(|| {
self.transact(|| {
Spi::connect(|client| {
pgrx::notice!("JSPG_SQL: {}", sql);
match client.select(sql, Some(args_with_oid.len() as i64), &args_with_oid) {
@ -41,11 +96,6 @@ impl DatabaseExecutor for SpiExecutor {
}
})
})
.catch_others(|cause| {
pgrx::warning!("JSPG Caught Native Postgres Error: {:?}", cause);
Err(format!("{:?}", cause))
})
.execute()
}
fn execute(&self, sql: &str, args: Option<&[Value]>) -> Result<(), String> {
@ -60,7 +110,7 @@ impl DatabaseExecutor for SpiExecutor {
}
}
pgrx::PgTryBuilder::new(|| {
self.transact(|| {
Spi::connect_mut(|client| {
pgrx::notice!("JSPG_SQL: {}", sql);
match client.update(sql, Some(args_with_oid.len() as i64), &args_with_oid) {
@ -69,44 +119,43 @@ impl DatabaseExecutor for SpiExecutor {
}
})
})
.catch_others(|cause| {
pgrx::warning!("JSPG Caught Native Postgres Error: {:?}", cause);
Err(format!("{:?}", cause))
})
.execute()
}
fn auth_user_id(&self) -> Result<String, String> {
Spi::connect(|client| {
let mut tup_table = client
.select(
"SELECT COALESCE(current_setting('auth.user_id', true), 'ffffffff-ffff-ffff-ffff-ffffffffffff')",
None,
&[],
)
.map_err(|e| format!("SPI Select Error: {}", e))?;
self.transact(|| {
Spi::connect(|client| {
let mut tup_table = client
.select(
"SELECT COALESCE(current_setting('auth.user_id', true), 'ffffffff-ffff-ffff-ffff-ffffffffffff')",
None,
&[],
)
.map_err(|e| format!("SPI Select Error: {}", e))?;
let row = tup_table
.next()
.ok_or("No user id setting returned from context".to_string())?;
let user_id: Option<String> = row.get(1).map_err(|e| e.to_string())?;
let row = tup_table
.next()
.ok_or("No user id setting returned from context".to_string())?;
let user_id: Option<String> = row.get(1).map_err(|e| e.to_string())?;
user_id.ok_or("Missing user_id".to_string())
user_id.ok_or("Missing user_id".to_string())
})
})
}
fn timestamp(&self) -> Result<String, String> {
Spi::connect(|client| {
let mut tup_table = client
.select("SELECT clock_timestamp()::text", None, &[])
.map_err(|e| format!("SPI Select Error: {}", e))?;
self.transact(|| {
Spi::connect(|client| {
let mut tup_table = client
.select("SELECT clock_timestamp()::text", None, &[])
.map_err(|e| format!("SPI Select Error: {}", e))?;
let row = tup_table
.next()
.ok_or("No clock timestamp returned".to_string())?;
let timestamp: Option<String> = row.get(1).map_err(|e| e.to_string())?;
let row = tup_table
.next()
.ok_or("No clock timestamp returned".to_string())?;
let timestamp: Option<String> = row.get(1).map_err(|e| e.to_string())?;
timestamp.ok_or("Missing timestamp".to_string())
timestamp.ok_or("Missing timestamp".to_string())
})
})
}
}