diff --git a/src/lib.rs b/src/lib.rs index 83a2ff6..d305b15 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2018,7 +2018,8 @@ mod tests { WITH _mark AS ( SELECT pg_catalog.set_config( 'df.non_future_helper', - 'df.await_instance' || E'\\n' || pg_catalog.statement_timestamp()::text, + 'df.await_instance' || E'\n' || + ((extract(epoch FROM pg_catalog.statement_timestamp()) - 946684800) * 1000000)::bigint::text, true ) ) diff --git a/src/types.rs b/src/types.rs index 7280437..86290d0 100644 --- a/src/types.rs +++ b/src/types.rs @@ -10,7 +10,7 @@ use cron::Schedule as CronSchedule; use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC}; use serde::{Deserialize, Serialize}; use std::borrow::Cow; -use std::ffi::CString; +use std::ffi::{CStr, CString}; use std::str::FromStr; use std::sync::{Arc, OnceLock}; use std::time::Duration; @@ -1149,11 +1149,22 @@ pub const VALID_NODE_TYPES: &[&str] = &[ const NON_FUTURE_HELPER_GUC: &str = "df.non_future_helper"; pub fn mark_non_future_helper_call(function_name: &str) { - if let Err(e) = Spi::run_with_args( - "SELECT pg_catalog.set_config('df.non_future_helper', $1 || E'\n' || pg_catalog.statement_timestamp()::text, true)", - &[function_name.into()], - ) { - pgrx::error!("Failed to mark helper call: {:?}", e); + let name = CString::new(NON_FUTURE_HELPER_GUC).expect("GUC name must not contain NUL bytes"); + let statement_timestamp = unsafe { pgrx::pg_sys::GetCurrentStatementStartTimestamp() }; + let marker = CString::new(format!("{}\n{}", function_name, statement_timestamp)) + .expect("helper name must not contain NUL bytes"); + + unsafe { + pgrx::pg_sys::set_config_option( + name.as_ptr(), + marker.as_ptr(), + pgrx::pg_sys::GucContext::PGC_USERSET, + pgrx::pg_sys::GucSource::PGC_S_SESSION, + pgrx::pg_sys::GucAction::GUC_ACTION_LOCAL, + true, + pgrx::PgLogLevel::ERROR as i32, + false, + ); } } @@ -1188,17 +1199,15 @@ impl Durofut { return None; } - let marker = Spi::get_one::(&format!( - "SELECT pg_catalog.current_setting('{}', true)", - NON_FUTURE_HELPER_GUC - )) - .ok() - .flatten()?; + let name = + CString::new(NON_FUTURE_HELPER_GUC).expect("GUC name must not contain NUL bytes"); + let marker = unsafe { + let value = pgrx::pg_sys::GetConfigOption(name.as_ptr(), true, false); + (!value.is_null()).then(|| CStr::from_ptr(value).to_string_lossy().into_owned()) + }?; let (helper_name, marker_timestamp) = marker.split_once('\n')?; - let statement_timestamp = - Spi::get_one::("SELECT pg_catalog.statement_timestamp()::text") - .ok() - .flatten()?; + let marker_timestamp = marker_timestamp.parse::().ok()?; + let statement_timestamp = unsafe { pgrx::pg_sys::GetCurrentStatementStartTimestamp() }; (marker_timestamp == statement_timestamp).then(|| helper_name.to_string()) }