diff --git a/CLAUDE.md b/CLAUDE.md index a25019ed..fe702708 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -53,12 +53,13 @@ pub trait ForeignDataWrapper> { fn end_modify(&mut self) -> Result<(), E>; // Optional methods for aggregate pushdown - fn supported_aggregates(&self) -> Vec; - fn supports_group_by(&self) -> bool; + fn supported_aggregates() -> Vec; + fn supports_group_by() -> bool; fn begin_aggregate_scan(&mut self, aggregates: &[Aggregate], group_by: &[Column], quals: &[Qual], options: &HashMap) -> Result<(), E>; // Optional methods fn re_scan(&mut self) -> Result<(), E>; + // Called during planning, before any FDW instance exists — takes no `self`. fn get_rel_size(...) -> Result<(i64, i32), E>; fn import_foreign_schema(...) -> Result, E>; fn validator(options: Vec>, catalog: Option) -> Result<(), E>; @@ -315,10 +316,10 @@ Use `Qual::deparse()` to convert to SQL-like strings. ### Aggregate Pushdown -FDWs can push `COUNT`, `SUM`, `AVG`, `MIN`, `MAX` (with optional `GROUP BY`) down to the remote source by implementing three optional trait methods: +FDWs can push `COUNT`, `SUM`, `AVG`, `MIN`, `MAX` (with optional `GROUP BY`) down to the remote source by implementing three optional trait methods. `supported_aggregates`/`supports_group_by` are called during query planning, before any instance of the FDW exists — they take no `self` and must not depend on FDW-instance state: ```rust -fn supported_aggregates(&self) -> Vec { +fn supported_aggregates() -> Vec { vec![ AggregateKind::Count, AggregateKind::CountColumn, @@ -329,7 +330,7 @@ fn supported_aggregates(&self) -> Vec { ] } -fn supports_group_by(&self) -> bool { true } +fn supports_group_by() -> bool { true } fn begin_aggregate_scan( &mut self, diff --git a/docs/guides/query-pushdown.md b/docs/guides/query-pushdown.md index bc4589a4..db592b0a 100644 --- a/docs/guides/query-pushdown.md +++ b/docs/guides/query-pushdown.md @@ -78,11 +78,11 @@ The Wrappers framework supports pushing down these aggregate functions: FDW developers can enable aggregate pushdown by implementing these trait methods: ```rust -fn supported_aggregates(&self) -> Vec { +fn supported_aggregates() -> Vec { vec![AggregateKind::Count, AggregateKind::Sum, AggregateKind::Avg] } -fn supports_group_by(&self) -> bool { +fn supports_group_by() -> bool { true } diff --git a/supabase-wrappers/src/interface.rs b/supabase-wrappers/src/interface.rs index 10c01ec7..13f8ab05 100644 --- a/supabase-wrappers/src/interface.rs +++ b/supabase-wrappers/src/interface.rs @@ -551,6 +551,15 @@ pub struct Qual { pub value: Value, pub use_or: bool, pub param: Option, + + // Stores the address of the original const node this qual's value was decoded + // from, if any. This is only used during serialization/deserialization to + // smuggle the `value` field across the planning and execution phase boundaries + // in fdw_private. This ensures the Qual survives Postgres' plan-cache + // `copyObject` call correctly. It's a usize instead of a *mut pg_sys::Const + // to keep `Qual` `Send`, which is important for FDWs like ClickHouse that + // move quals across tokio task boundaries. + pub(crate) value_const: Option, } impl Qual { @@ -881,6 +890,12 @@ pub trait ForeignDataWrapper> { /// You can do any initalization in this function, like saving connection /// info or API url in an variable, but don't do heavy works like database /// connection or API call. + /// + /// Never called during query planning — [`get_rel_size`](Self::get_rel_size), + /// [`supported_aggregates`](Self::supported_aggregates) and + /// [`supports_group_by`](Self::supports_group_by) are the only planning-time + /// hooks, and none of them take a `self`. `new` only runs once per actual + /// execution (once per `EXECUTE` of a cached/prepared plan). fn new(server: ForeignServer) -> Result where Self: Sized; @@ -890,9 +905,12 @@ pub trait ForeignDataWrapper> { /// Return the expected number of rows and row size (in bytes) by the /// foreign table scan. /// + /// Called during query planning, before any instance of this FDW exists for the + /// query (planning never constructs one, see `new`'s docs) — implementations must + /// not depend on any FDW-instance state. + /// /// [See more details](https://www.postgresql.org/docs/current/fdw-callbacks.html#FDW-CALLBACKS-SCAN). fn get_rel_size( - &mut self, _quals: &[Qual], _columns: &[Column], _sorts: &[Sort], @@ -1018,10 +1036,13 @@ pub trait ForeignDataWrapper> { /// /// ## Examples /// + /// Called during query planning, before any instance of this FDW exists for the + /// query — implementations must not depend on any FDW-instance state. + /// /// ```rust,ignore /// use supabase_wrappers::prelude::*; /// - /// fn supported_aggregates(&self) -> Vec { + /// fn supported_aggregates() -> Vec { /// vec![ /// AggregateKind::Count, /// AggregateKind::CountColumn, @@ -1032,7 +1053,7 @@ pub trait ForeignDataWrapper> { /// ] /// } /// ``` - fn supported_aggregates(&self) -> Vec { + fn supported_aggregates() -> Vec { vec![] } @@ -1043,14 +1064,17 @@ pub trait ForeignDataWrapper> { /// /// When `true`, GROUP BY columns will be passed to [`begin_aggregate_scan`](Self::begin_aggregate_scan). /// + /// Called during query planning, before any instance of this FDW exists for the + /// query — implementations must not depend on any FDW-instance state. + /// /// ## Examples /// /// ```rust,ignore - /// fn supports_group_by(&self) -> bool { + /// fn supports_group_by() -> bool { /// true /// } /// ``` - fn supports_group_by(&self) -> bool { + fn supports_group_by() -> bool { false } @@ -1184,7 +1208,11 @@ pub trait ForeignDataWrapper> { Ok(Vec::new()) } - /// Returns a FdwRoutine for the FDW + /// The handler function for all foreign data wrappers. + /// + /// The [`FdwRoutine`] is the same as the `fdw_handler` pseudo-type mentioned in the + /// [Postgres documentation](https://www.postgresql.org/docs/current/fdw-functions.html). + /// This is the entry point of a foreign table query: the first callback called by Postgres. /// /// Not to be used directly, use [`wrappers_fdw`](crate::wrappers_fdw) macro instead. fn fdw_routine() -> FdwRoutine diff --git a/supabase-wrappers/src/qual.rs b/supabase-wrappers/src/qual.rs index 429240da..5ce5d213 100644 --- a/supabase-wrappers/src/qual.rs +++ b/supabase-wrappers/src/qual.rs @@ -175,7 +175,6 @@ pub(crate) unsafe fn unnest_clause(node: *mut pg_sys::Node) -> *mut pg_sys::Node } pub(crate) unsafe fn extract_from_op_expr( - _root: *mut pg_sys::PlannerInfo, baserel_id: pg_sys::Oid, baserel_ids: pg_sys::Relids, expr: *mut pg_sys::OpExpr, @@ -216,7 +215,7 @@ pub(crate) unsafe fn extract_from_op_expr( { let field = pg_sys::get_attname(baserel_id, (*left).varattno, false); - let (value, param) = if is_a(right, pg_sys::NodeTag::T_Const) { + let (value, param, value_const) = if is_a(right, pg_sys::NodeTag::T_Const) { let right = right as *mut pg_sys::Const; ( Cell::from_polymorphic_datum( @@ -225,6 +224,7 @@ pub(crate) unsafe fn extract_from_op_expr( (*right).consttype, ), None, + Some(right as usize), ) } else if is_a(right, pg_sys::NodeTag::T_Param) { // add a dummy value if this is query parameter, the actual value @@ -244,9 +244,9 @@ pub(crate) unsafe fn extract_from_op_expr( expr_state: ptr::null_mut(), }, }; - (Some(Cell::I64(0)), Some(param)) + (Some(Cell::I64(0)), Some(param), None) } else { - (None, None) + (None, None, None) }; if let Some(value) = value { @@ -256,6 +256,7 @@ pub(crate) unsafe fn extract_from_op_expr( value: Value::Cell(value), use_or: false, param, + value_const, }; return Some(qual); } @@ -296,6 +297,7 @@ pub(crate) unsafe fn extract_from_null_test( value: Value::Cell(Cell::String("null".to_string())), use_or: false, param: None, + value_const: None, }; Some(qual) @@ -303,7 +305,6 @@ pub(crate) unsafe fn extract_from_null_test( } pub(crate) unsafe fn extract_from_scalar_array_op_expr( - _root: *mut pg_sys::PlannerInfo, baserel_id: pg_sys::Oid, baserel_ids: pg_sys::Relids, expr: *mut pg_sys::ScalarArrayOpExpr, @@ -347,6 +348,7 @@ pub(crate) unsafe fn extract_from_scalar_array_op_expr( value: Value::Array(value), use_or: (*expr).useOr, param: None, + value_const: Some(right as usize), }; return Some(qual); } @@ -364,7 +366,6 @@ pub(crate) unsafe fn extract_from_scalar_array_op_expr( } pub(crate) unsafe fn extract_from_var( - _root: *mut pg_sys::PlannerInfo, baserel_id: pg_sys::Oid, baserel_ids: pg_sys::Relids, var: *mut pg_sys::Var, @@ -385,6 +386,7 @@ pub(crate) unsafe fn extract_from_var( value: Value::Cell(Cell::Bool(true)), use_or: false, param: None, + value_const: None, }; Some(qual) @@ -392,7 +394,6 @@ pub(crate) unsafe fn extract_from_var( } pub(crate) unsafe fn extract_from_bool_expr( - _root: *mut pg_sys::PlannerInfo, baserel_id: pg_sys::Oid, baserel_ids: pg_sys::Relids, expr: *mut pg_sys::BoolExpr, @@ -420,6 +421,7 @@ pub(crate) unsafe fn extract_from_bool_expr( value: Value::Cell(Cell::Bool(false)), use_or: false, param: None, + value_const: None, }; return Some(qual); @@ -456,6 +458,7 @@ pub(crate) unsafe fn extract_from_boolean_test( value: Value::Cell(Cell::Bool(value)), use_or: false, param: None, + value_const: None, }; Some(qual) @@ -463,7 +466,6 @@ pub(crate) unsafe fn extract_from_boolean_test( } pub(crate) unsafe fn extract_quals( - root: *mut pg_sys::PlannerInfo, baserel: *mut pg_sys::RelOptInfo, baserel_id: pg_sys::Oid, ) -> Vec { @@ -477,20 +479,15 @@ pub(crate) unsafe fn extract_quals( for cond in conds.iter() { let expr = (*(*cond as *mut pg_sys::RestrictInfo)).clause as *mut pg_sys::Node; let extracted = if is_a(expr, pg_sys::NodeTag::T_OpExpr) { - extract_from_op_expr(root, baserel_id, (*baserel).relids, expr as _) + extract_from_op_expr(baserel_id, (*baserel).relids, expr as _) } else if is_a(expr, pg_sys::NodeTag::T_NullTest) { extract_from_null_test(baserel_id, expr as _) } else if is_a(expr, pg_sys::NodeTag::T_ScalarArrayOpExpr) { - extract_from_scalar_array_op_expr( - root, - baserel_id, - (*baserel).relids, - expr as _, - ) + extract_from_scalar_array_op_expr(baserel_id, (*baserel).relids, expr as _) } else if is_a(expr, pg_sys::NodeTag::T_Var) { - extract_from_var(root, baserel_id, (*baserel).relids, expr as _) + extract_from_var(baserel_id, (*baserel).relids, expr as _) } else if is_a(expr, pg_sys::NodeTag::T_BoolExpr) { - extract_from_bool_expr(root, baserel_id, (*baserel).relids, expr as _) + extract_from_bool_expr(baserel_id, (*baserel).relids, expr as _) } else if is_a(expr, pg_sys::NodeTag::T_BooleanTest) { extract_from_boolean_test(baserel_id, expr as _) } else { diff --git a/supabase-wrappers/src/scan.rs b/supabase-wrappers/src/scan.rs index 436c7d7d..2aacb245 100644 --- a/supabase-wrappers/src/scan.rs +++ b/supabase-wrappers/src/scan.rs @@ -1,19 +1,26 @@ use pgrx::FromDatum; use pgrx::{ IntoDatum, PgSqlErrorCode, debug2, + list::List, + memcx::MemCx, memcxt::PgMemoryContexts, pg_sys::{Datum, MemoryContext, MemoryContextData, Oid, ParamKind}, prelude::*, }; use std::collections::HashMap; +use std::ffi::c_void; use std::marker::PhantomData; +use std::mem; +use std::sync::Mutex; use pgrx::pg_sys::panic::ErrorReport; use std::os::raw::c_int; use std::ptr; use crate::instance; -use crate::interface::{Aggregate, Cell, Column, Limit, Qual, Row, Sort, Value}; +use crate::interface::{ + Aggregate, AggregateKind, Cell, Column, ExprEval, Limit, Param, Qual, Row, Sort, Value, +}; use crate::limit::*; use crate::memctx; use crate::options::options_to_hashmap; @@ -21,10 +28,18 @@ use crate::polyfill; use crate::prelude::ForeignDataWrapper; use crate::qual::*; use crate::sort::*; -use crate::utils::{self, ReportableError, SerdeList, report_error}; +use crate::utils::{self, ReportableError, report_error}; // Fdw private state for scan pub(crate) struct FdwState, W: ForeignDataWrapper> { + // The base relation's foreign table Oid, captured once during + // `get_foreign_rel_size` (always called for the base rel, so always valid). + // `get_foreign_plan` must use this rather than its own `foreigntableid` + // parameter: for an aggregate-pushdown plan, `baserel` there is the upper + // (GROUP_AGG) relation, and Postgres passes `InvalidOid` in that case since + // an upper rel isn't tied to a single base relation. + pub(crate) foreigntableid: Oid, + // foreign data wrapper instance pub(crate) instance: Option, @@ -61,9 +76,15 @@ pub(crate) struct FdwState, W: ForeignDataWrapper> { } impl, W: ForeignDataWrapper> FdwState { + // Used only for planning (`get_foreign_rel_size`). `get_rel_size`, + // `supported_aggregates` and `supports_group_by` are the only planning-time + // trait hooks, and none of them take a `self`, so planning never needs a + // live FDW instance — leaving `instance: None` here means the (potentially + // expensive) `W::new()` only ever runs once per actual execution. unsafe fn new(foreigntableid: Oid, tmp_ctx: MemoryContext) -> Self { Self { - instance: Some(unsafe { instance::create_fdw_instance_from_table_id(foreigntableid) }), + foreigntableid, + instance: None, quals: Vec::new(), tgts: Vec::new(), sorts: Vec::new(), @@ -82,17 +103,13 @@ impl, W: ForeignDataWrapper> FdwState { #[inline] fn get_rel_size(&mut self) -> Result<(i64, i32), E> { - if let Some(ref mut instance) = self.instance { - instance.get_rel_size( - &self.quals, - &self.tgts, - &self.sorts, - &self.limit, - &self.opts, - ) - } else { - Ok((0, 0)) - } + W::get_rel_size( + &self.quals, + &self.tgts, + &self.sorts, + &self.limit, + &self.opts, + ) } #[inline] @@ -152,8 +169,6 @@ impl, W: ForeignDataWrapper> FdwState { } } -impl, W: ForeignDataWrapper> utils::SerdeList for FdwState {} - impl, W: ForeignDataWrapper> Drop for FdwState { fn drop(&mut self) { // drop foreign data wrapper instance @@ -175,6 +190,615 @@ unsafe fn drop_fdw_state, W: ForeignDataWrapper>( drop(boxed_fdw_state); } +/// This struct is a serializable state of the planning time data needed to +/// rebuild [`FdwState`] in the execution phase. +/// +/// Unlike [`FdwState`] which owns a live FDW instance, a Postgres MemoryContext, +/// and per-scan row buffers, this struct holds only plain data. This struct will +/// be serialized as a [`pg_sys::List`] of [`pg_sys::Const`] nodes so that when +/// Postgres calls `copyObject` on it at the end of the plan phase (after the +/// function call [`get_foreign_plan`]) it is deep copied correctly and rebuilt +/// successfully at the beginning of the [`begin_foreign_scan`] function. +struct FdwScanPrivate { + foreigntableid: Oid, + quals: Vec, + tgts: Vec, + sorts: Vec, + limit: Option, + aggregates: Vec, + group_by: Vec, +} + +/// How a `Qual::value` is encoded in [`FdwScanPrivate`]'s serialized list. `ScalarConst`/ +/// `ArrayConst` embed the original `pg_sys::Const` node (see [`Qual::value_const`]) so +/// `copyObject` deep-copies it with the correct `consttype`; `Bool` and `Placeholder` +/// have no source `Const` node to preserve (see `push_qual`/`read_qual`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum QualValueMode { + Bool = 0, + Placeholder = 1, + ScalarConst = 2, + ArrayConst = 3, +} + +impl QualValueMode { + fn from_i32(val: i32) -> Option { + match val { + 0 => Some(Self::Bool), + 1 => Some(Self::Placeholder), + 2 => Some(Self::ScalarConst), + 3 => Some(Self::ArrayConst), + _ => None, + } + } +} + +impl FdwScanPrivate { + unsafe fn serialize_to_list(&self) -> *mut pg_sys::List { + unsafe { + pgrx::memcx::current_context(|mcx| { + let mut ret = List::<*mut c_void>::Nil; + Self::push_oid(&mut ret, mcx, self.foreigntableid); + Self::push_quals(&mut ret, mcx, &self.quals); + Self::push_columns(&mut ret, mcx, &self.tgts); + Self::push_sorts(&mut ret, mcx, &self.sorts); + Self::push_limit(&mut ret, mcx, &self.limit); + Self::push_aggregates(&mut ret, mcx, &self.aggregates); + Self::push_columns(&mut ret, mcx, &self.group_by); + ret.into_ptr() + }) + } + } + + unsafe fn deserialize_from_list(list: *mut pg_sys::List) -> Option { + unsafe { + pgrx::memcx::current_context(|mcx| { + let list = List::<*mut c_void>::downcast_ptr_in_memcx(list, mcx)?; + let mut idx = 0usize; + + let foreigntableid = Self::read_oid(&list, &mut idx)?; + let quals = Self::read_quals(&list, &mut idx)?; + let tgts = Self::read_columns(&list, &mut idx)?; + let sorts = Self::read_sorts(&list, &mut idx)?; + let limit = Self::read_limit(&list, &mut idx)?; + let aggregates = Self::read_aggregates(&list, &mut idx)?; + let group_by = Self::read_columns(&list, &mut idx)?; + + Some(FdwScanPrivate { + foreigntableid, + quals, + tgts, + sorts, + limit, + aggregates, + group_by, + }) + }) + } + } + + unsafe fn push_i32<'cx>(list: &mut List<'cx, *mut c_void>, mcx: &'cx MemCx<'_>, val: i32) { + unsafe { + let cst = pg_sys::makeConst( + pg_sys::INT4OID, + -1, + pg_sys::InvalidOid, + 4, + val.into_datum().unwrap(), + false, + true, + ); + list.unstable_push_in_context(cst as _, mcx); + } + } + + unsafe fn push_i64<'cx>(list: &mut List<'cx, *mut c_void>, mcx: &'cx MemCx<'_>, val: i64) { + unsafe { + let cst = pg_sys::makeConst( + pg_sys::INT8OID, + -1, + pg_sys::InvalidOid, + 8, + val.into_datum().unwrap(), + false, + true, + ); + list.unstable_push_in_context(cst as _, mcx); + } + } + + unsafe fn push_bool<'cx>(list: &mut List<'cx, *mut c_void>, mcx: &'cx MemCx<'_>, val: bool) { + unsafe { + let cst = pg_sys::makeConst( + pg_sys::BOOLOID, + -1, + pg_sys::InvalidOid, + 1, + val.into_datum().unwrap(), + false, + true, + ); + list.unstable_push_in_context(cst as _, mcx); + } + } + + unsafe fn push_text<'cx>(list: &mut List<'cx, *mut c_void>, mcx: &'cx MemCx<'_>, val: &str) { + unsafe { + let cst = pg_sys::makeConst( + pg_sys::TEXTOID, + -1, + pg_sys::InvalidOid, + -1, + val.to_string().into_datum().unwrap(), + false, + false, + ); + list.unstable_push_in_context(cst as _, mcx); + } + } + + unsafe fn push_oid<'cx>(list: &mut List<'cx, *mut c_void>, mcx: &'cx MemCx<'_>, val: Oid) { + unsafe { Self::push_i32(list, mcx, val.to_u32() as i32) }; + } + + // Reads the raw `Const` at the current cursor position and advances the cursor. + unsafe fn read_const(list: &List<*mut c_void>, idx: &mut usize) -> Option { + let cst_ptr = *list.get(*idx)? as *mut pg_sys::Const; + *idx += 1; + Some(unsafe { *cst_ptr }) + } + + unsafe fn read_i32(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let cst = Self::read_const(list, idx)?; + i32::from_datum(cst.constvalue, cst.constisnull) + } + } + + unsafe fn read_i64(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let cst = Self::read_const(list, idx)?; + i64::from_datum(cst.constvalue, cst.constisnull) + } + } + + unsafe fn read_bool(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let cst = Self::read_const(list, idx)?; + bool::from_datum(cst.constvalue, cst.constisnull) + } + } + + unsafe fn read_text(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let cst = Self::read_const(list, idx)?; + String::from_datum(cst.constvalue, cst.constisnull) + } + } + + unsafe fn read_oid(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { Self::read_i32(list, idx) }.map(|v| Oid::from(v as u32)) + } + + unsafe fn push_column<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + col: &Column, + ) { + unsafe { + Self::push_text(list, mcx, &col.name); + // usize to i32 cast is safe as Postgres has a maximum of 1600 columns + Self::push_i32(list, mcx, col.num as i32); + Self::push_oid(list, mcx, col.type_oid); + } + } + + unsafe fn read_column(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let name = Self::read_text(list, idx)?; + let num = Self::read_i32(list, idx)? as usize; + let type_oid = Self::read_oid(list, idx)?; + Some(Column { + name, + num, + type_oid, + }) + } + } + + unsafe fn push_columns<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + cols: &[Column], + ) { + unsafe { + Self::push_i32(list, mcx, cols.len() as i32); + for col in cols { + Self::push_column(list, mcx, col); + } + } + } + + unsafe fn read_columns(list: &List<*mut c_void>, idx: &mut usize) -> Option> { + unsafe { + let count = Self::read_i32(list, idx)? as usize; + let mut cols = Vec::with_capacity(count); + for _ in 0..count { + cols.push(Self::read_column(list, idx)?); + } + Some(cols) + } + } + + unsafe fn push_sort<'cx>(list: &mut List<'cx, *mut c_void>, mcx: &'cx MemCx<'_>, sort: &Sort) { + unsafe { + Self::push_text(list, mcx, &sort.field); + // usize to i32 cast is safe field_no is also bound by Postgres maximum number of columns(1600) + Self::push_i32(list, mcx, sort.field_no as i32); + Self::push_bool(list, mcx, sort.reversed); + Self::push_bool(list, mcx, sort.nulls_first); + Self::push_bool(list, mcx, sort.collate.is_some()); + if let Some(collate) = &sort.collate { + Self::push_text(list, mcx, collate); + } + } + } + + unsafe fn read_sort(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let field = Self::read_text(list, idx)?; + let field_no = Self::read_i32(list, idx)? as usize; + let reversed = Self::read_bool(list, idx)?; + let nulls_first = Self::read_bool(list, idx)?; + let has_collate = Self::read_bool(list, idx)?; + let collate = if has_collate { + Some(Self::read_text(list, idx)?) + } else { + None + }; + + Some(Sort { + field, + field_no, + reversed, + nulls_first, + collate, + }) + } + } + + unsafe fn push_sorts<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + sorts: &[Sort], + ) { + unsafe { + Self::push_i32(list, mcx, sorts.len() as i32); + for sort in sorts { + Self::push_sort(list, mcx, sort); + } + } + } + + unsafe fn read_sorts(list: &List<*mut c_void>, idx: &mut usize) -> Option> { + unsafe { + let count = Self::read_i32(list, idx)? as usize; + let mut sorts = Vec::with_capacity(count); + for _ in 0..count { + sorts.push(Self::read_sort(list, idx)?); + } + Some(sorts) + } + } + + unsafe fn push_limit<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + limit: &Option, + ) { + unsafe { + Self::push_bool(list, mcx, limit.is_some()); + if let Some(limit) = limit { + Self::push_i64(list, mcx, limit.count); + Self::push_i64(list, mcx, limit.offset); + } + } + } + + unsafe fn read_limit(list: &List<*mut c_void>, idx: &mut usize) -> Option> { + unsafe { + let has_limit = Self::read_bool(list, idx)?; + if has_limit { + let count = Self::read_i64(list, idx)?; + let offset = Self::read_i64(list, idx)?; + Some(Some(Limit { count, offset })) + } else { + Some(None) + } + } + } + + fn aggregate_kind_to_i32(kind: AggregateKind) -> i32 { + match kind { + AggregateKind::Count => 0, + AggregateKind::CountColumn => 1, + AggregateKind::Sum => 2, + AggregateKind::Avg => 3, + AggregateKind::Min => 4, + AggregateKind::Max => 5, + } + } + + fn aggregate_kind_from_i32(val: i32) -> Option { + match val { + 0 => Some(AggregateKind::Count), + 1 => Some(AggregateKind::CountColumn), + 2 => Some(AggregateKind::Sum), + 3 => Some(AggregateKind::Avg), + 4 => Some(AggregateKind::Min), + 5 => Some(AggregateKind::Max), + _ => None, + } + } + + unsafe fn push_aggregate<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + agg: &Aggregate, + ) { + unsafe { + Self::push_i32(list, mcx, Self::aggregate_kind_to_i32(agg.kind)); + Self::push_bool(list, mcx, agg.column.is_some()); + if let Some(col) = &agg.column { + Self::push_column(list, mcx, col); + } + Self::push_bool(list, mcx, agg.distinct); + Self::push_text(list, mcx, &agg.alias); + Self::push_oid(list, mcx, agg.type_oid); + } + } + + unsafe fn read_aggregate(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let kind = Self::aggregate_kind_from_i32(Self::read_i32(list, idx)?)?; + let has_column = Self::read_bool(list, idx)?; + let column = if has_column { + Some(Self::read_column(list, idx)?) + } else { + None + }; + let distinct = Self::read_bool(list, idx)?; + let alias = Self::read_text(list, idx)?; + let type_oid = Self::read_oid(list, idx)?; + Some(Aggregate { + kind, + column, + distinct, + alias, + type_oid, + }) + } + } + + unsafe fn push_aggregates<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + aggregates: &[Aggregate], + ) { + unsafe { + Self::push_i32(list, mcx, aggregates.len() as i32); + for agg in aggregates { + Self::push_aggregate(list, mcx, agg); + } + } + } + + unsafe fn read_aggregates(list: &List<*mut c_void>, idx: &mut usize) -> Option> { + unsafe { + let count = Self::read_i32(list, idx)? as usize; + let mut aggregates = Vec::with_capacity(count); + for _ in 0..count { + aggregates.push(Self::read_aggregate(list, idx)?); + } + Some(aggregates) + } + } + + unsafe fn push_qual<'cx>(list: &mut List<'cx, *mut c_void>, mcx: &'cx MemCx<'_>, qual: &Qual) { + unsafe { + Self::push_text(list, mcx, &qual.field); + Self::push_text(list, mcx, &qual.operator); + Self::push_bool(list, mcx, qual.use_or); + + match qual.value_const { + Some(addr) => { + let mode = if matches!(qual.value, Value::Array(_)) { + QualValueMode::ArrayConst + } else { + QualValueMode::ScalarConst + }; + Self::push_i32(list, mcx, mode as i32); + list.unstable_push_in_context(addr as *mut c_void, mcx); + } + None => match &qual.value { + Value::Cell(Cell::Bool(b)) => { + Self::push_i32(list, mcx, QualValueMode::Bool as i32); + Self::push_bool(list, mcx, *b); + } + _ => { + Self::push_i32(list, mcx, QualValueMode::Placeholder as i32); + } + }, + } + + Self::push_param(list, mcx, &qual.param); + } + } + + unsafe fn push_param<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + param: &Option, + ) { + unsafe { + Self::push_bool(list, mcx, param.is_some()); + if let Some(param) = param { + Self::push_i32(list, mcx, param.kind as i32); + Self::push_i32(list, mcx, param.id as i32); + Self::push_oid(list, mcx, param.type_oid); + } + } + } + + unsafe fn push_quals<'cx>( + list: &mut List<'cx, *mut c_void>, + mcx: &'cx MemCx<'_>, + quals: &[Qual], + ) { + unsafe { + Self::push_i32(list, mcx, quals.len() as i32); + for qual in quals { + Self::push_qual(list, mcx, qual); + } + } + } + + unsafe fn read_qual(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let field = Self::read_text(list, idx)?; + let operator = Self::read_text(list, idx)?; + let use_or = Self::read_bool(list, idx)?; + + let mode = QualValueMode::from_i32(Self::read_i32(list, idx)?)?; + let value = match mode { + QualValueMode::Bool => Value::Cell(Cell::Bool(Self::read_bool(list, idx)?)), + QualValueMode::Placeholder => Value::Cell(Cell::String("null".to_string())), + QualValueMode::ScalarConst => { + let cst = Self::read_const(list, idx)?; + Value::Cell(Cell::from_polymorphic_datum( + cst.constvalue, + cst.constisnull, + cst.consttype, + )?) + } + QualValueMode::ArrayConst => { + let cst = Self::read_const(list, idx)?; + Value::Array(form_array_from_datum( + cst.constvalue, + cst.constisnull, + cst.consttype, + )?) + } + }; + + let param = Self::read_param(list, idx); + + Some(Qual { + field, + operator, + value, + use_or, + param, + value_const: None, + }) + } + } + + unsafe fn read_param(list: &List<*mut c_void>, idx: &mut usize) -> Option { + unsafe { + let has_param = Self::read_bool(list, idx)?; + if has_param { + let kind = Self::read_i32(list, idx)? as pg_sys::ParamKind::Type; + let id = Self::read_i32(list, idx)? as usize; + let type_oid = Self::read_oid(list, idx)?; + Some(Param { + kind, + id, + type_oid, + eval_value: Mutex::new(None).into(), + expr_eval: ExprEval { + expr: ptr::null_mut(), + expr_state: ptr::null_mut(), + }, + }) + } else { + None + } + } + } + + unsafe fn read_quals(list: &List<*mut c_void>, idx: &mut usize) -> Option> { + unsafe { + let count = Self::read_i32(list, idx)? as usize; + let mut quals = Vec::with_capacity(count); + for _ in 0..count { + quals.push(Self::read_qual(list, idx)?); + } + Some(quals) + } + } +} + +impl, W: ForeignDataWrapper> FdwState { + /// Deserialize [`FdwState`] from a [`FdwScanPrivate`] struct. + unsafe fn from_scan_private(private: FdwScanPrivate, tmp_ctx: MemoryContext) -> Self { + unsafe { + let foreigntableid = private.foreigntableid; + let instance = instance::create_fdw_instance_from_table_id(foreigntableid); + + let ftable = pg_sys::GetForeignTable(foreigntableid); + let mut opts = options_to_hashmap((*ftable).options).report_unwrap(); + opts.insert( + "wrappers.fserver_oid".into(), + (*ftable).serverid.to_u32().to_string(), + ); + opts.insert( + "wrappers.ftable_oid".into(), + (*ftable).relid.to_u32().to_string(), + ); + + let mut quals = private.quals; + + // Reallocate the `pg_sys::ParamKind::PARAM_EXEC` node in the `tmp_ctx` + // memory context. + PgMemoryContexts::For(tmp_ctx).switch_to(|_| { + for qual in &mut quals { + if let Some(param) = &mut qual.param + && param.kind == pg_sys::ParamKind::PARAM_EXEC + { + let mut node = PgBox::::alloc_node(pg_sys::NodeTag::T_Param); + node.paramkind = param.kind; + node.paramid = param.id as _; + node.paramtype = param.type_oid; + node.paramtypmod = -1; + node.paramcollid = pg_sys::InvalidOid; + node.location = -1; + param.expr_eval.expr = node.into_pg() as _; + } + } + }); + + Self { + foreigntableid, + instance: Some(instance), + quals, + tgts: private.tgts, + sorts: private.sorts, + limit: private.limit, + opts, + aggregates: private.aggregates, + group_by: private.group_by, + tmp_ctx, + values: Vec::new(), + nulls: Vec::new(), + row: Row::new(), + param_fingerprint: String::new(), + _phantom: PhantomData, + } + } + } +} + #[pg_guard] pub(super) extern "C-unwind" fn get_foreign_rel_size< E: Into, @@ -195,7 +819,7 @@ pub(super) extern "C-unwind" fn get_foreign_rel_size< PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| { // extract qual list - state.quals = extract_quals(root, baserel, foreigntableid); + state.quals = extract_quals(baserel, foreigntableid); // extract target column list from target and restriction expression state.tgts = utils::extract_target_columns(root, baserel); @@ -282,6 +906,8 @@ pub(super) extern "C-unwind" fn get_foreign_paths< pub(super) extern "C-unwind" fn get_foreign_plan, W: ForeignDataWrapper>( _root: *mut pg_sys::PlannerInfo, baserel: *mut pg_sys::RelOptInfo, + // Not `state.foreigntableid`'s source: unreliable (`InvalidOid`) for + // aggregate-pushdown (upper-rel) plans — see the comment below. _foreigntableid: pg_sys::Oid, _best_path: *mut pg_sys::ForeignPath, tlist: *mut pg_sys::List, @@ -384,14 +1010,31 @@ pub(super) extern "C-unwind" fn get_foreign_plan, W: Foreig (tlist, ptr::null_mut()) }; - // 'serialize' state to list, basically what we're doing here is to store - // the state pointer as an integer constant in the list, so it can be - // `deserialized` when executing the plan later. - // Note that the state itself is not serialized to any memory contexts, - // it just sits in Rust managed Box'ed memory and will be dropped when - // end_foreign_scan() is called. - let fdw_private = - PgMemoryContexts::For(state.tmp_ctx).switch_to(|_| FdwState::serialize_to_list(state)); + // It is critical that the data we pass in `fdw_private` be deep copyable + // via a Postgres `copyObject` call. Since `get_foreign_plan` is the last + // callback of the plan phase, Postgres needs to potentially be able to + // cache the plan and run the scan phase repeatedly using this cached plan. + // When Postgres runs the scan phase it calls `copyObject` on the plan (including + // `fdw_private`) before passing it to the scan phase's `begin_foreign_scan` + // callback where this state will be reconstituted. + // Use `state.foreigntableid` (captured for the base rel during + // `get_foreign_rel_size`), not this callback's own `foreigntableid` + // parameter: for an aggregate-pushdown plan, `baserel` here is the + // upper (GROUP_AGG) relation and Postgres passes `InvalidOid` for it. + let private = FdwScanPrivate { + foreigntableid: state.foreigntableid, + quals: mem::take(&mut state.quals), + tgts: mem::take(&mut state.tgts), + sorts: mem::take(&mut state.sorts), + limit: state.limit.take(), + aggregates: mem::take(&mut state.aggregates), + group_by: mem::take(&mut state.group_by), + }; + let fdw_private = private.serialize_to_list(); + + // Drop the state struct because its values have been serialized into + // `fdw_private` and it is no longer needed. + drop_fdw_state(state.as_ptr()); pg_sys::make_foreignscan( final_tlist, @@ -551,8 +1194,22 @@ pub(super) extern "C-unwind" fn begin_foreign_scan< unsafe { let scan_state = (*node).ss; let plan = scan_state.ps.plan as *mut pg_sys::ForeignScan; - let mut state = FdwState::::deserialize_from_list((*plan).fdw_private as _); - assert!(!state.is_null()); + + let Some(private) = FdwScanPrivate::deserialize_from_list((*plan).fdw_private as _) else { + report_error( + PgSqlErrorCode::ERRCODE_FDW_ERROR, + "invalid fdw_private data in begin_foreign_scan", + ); + return; + }; + + // Rebuild the scan state again from the serialized `FdwScanPrivate` afresh each time + // `begin_foreign_scan` is called to avoid state struct lifetime issues. The plan phase + // might have cached the plan, so we create a fresh copy in the scan phase. + let foreigntableid = private.foreigntableid; + let ctx_name = format!("Wrappers_scan_{}", foreigntableid.to_u32()); + let tmp_ctx = memctx::create_wrappers_memctx(&ctx_name); + let mut state = FdwState::::from_scan_private(private, tmp_ctx); // assign parameter values to qual assign_parameter_value(node, &mut state); @@ -566,11 +1223,7 @@ pub(super) extern "C-unwind" fn begin_foreign_scan< } else { state.begin_scan() }; - if result.is_err() { - drop_fdw_state(state.as_ptr()); - (*plan).fdw_private = ptr::null::>() as _; - result.report_unwrap(); - } + result.report_unwrap(); // For aggregate upper-rel scans, scanrelid=0 so ss_currentRelation is // NULL. Use the number of output columns from state.tgts instead. @@ -588,7 +1241,8 @@ pub(super) extern "C-unwind" fn begin_foreign_scan< state.nulls.extend_from_slice(&vec![true; natts]); } - (*node).fdw_state = state.into_pg() as _; + // This is leaked here but dropped in `end_foreign_scan` + (*node).fdw_state = Box::leak(Box::new(state)) as *mut FdwState as _; } } diff --git a/supabase-wrappers/src/upper.rs b/supabase-wrappers/src/upper.rs index 35856f7d..9467d66c 100644 --- a/supabase-wrappers/src/upper.rs +++ b/supabase-wrappers/src/upper.rs @@ -301,16 +301,10 @@ pub(super) extern "C-unwind" fn get_foreign_upper_paths< let mut state = PgBox::>::from_pg(fdw_private as _); // Check if FDW supports any aggregates - let supported = { - let Some(ref instance) = state.instance else { - return; - }; - let supported = instance.supported_aggregates(); - if supported.is_empty() { - return; - } - supported - }; + let supported = W::supported_aggregates(); + if supported.is_empty() { + return; + } // Extract aggregates from the query let aggregates = match extract_aggregates(root, output_rel, extra) { @@ -339,14 +333,9 @@ pub(super) extern "C-unwind" fn get_foreign_upper_paths< } // Check if GROUP BY is supported (if present) - if !group_by.is_empty() { - let Some(ref instance) = state.instance else { - return; - }; - if !instance.supports_group_by() { - debug2!("GROUP BY not supported, skipping pushdown"); - return; - } + if !group_by.is_empty() && !W::supports_group_by() { + debug2!("GROUP BY not supported, skipping pushdown"); + return; } // Store aggregates and group_by in the FdwState so they survive to diff --git a/supabase-wrappers/src/utils.rs b/supabase-wrappers/src/utils.rs index ae69cee7..215f5ed0 100644 --- a/supabase-wrappers/src/utils.rs +++ b/supabase-wrappers/src/utils.rs @@ -3,7 +3,6 @@ use crate::interface::{Cell, Column, Row}; use pgrx::{ - IntoDatum, list::List, pg_sys::panic::{ErrorReport, ErrorReportable}, spi::Spi, @@ -515,51 +514,6 @@ pub(super) unsafe fn extract_target_columns( } } -// trait for "serialize" and "deserialize" state from specified memory context, -// so that it is safe to be carried between the planning and the execution -pub(super) trait SerdeList { - unsafe fn serialize_to_list(state: PgBox) -> *mut pg_sys::List - where - Self: Sized, - { - unsafe { - memcx::current_context(|mcx| { - let mut ret = List::<*mut c_void>::Nil; - let val = state.into_pg() as i64; - let cst: *mut pg_sys::Const = pg_sys::makeConst( - pg_sys::INT8OID, - -1, - pg_sys::InvalidOid, - 8, - val.into_datum().unwrap(), - false, - true, - ); - ret.unstable_push_in_context(cst as _, mcx); - ret.into_ptr() - }) - } - } - - unsafe fn deserialize_from_list(list: *mut pg_sys::List) -> PgBox - where - Self: Sized, - { - unsafe { - memcx::current_context(|mcx| { - if let Some(list) = List::<*mut c_void>::downcast_ptr_in_memcx(list, mcx) - && let Some(cst) = list.get(0) - { - let cst = *(*cst as *mut pg_sys::Const); - let ptr = i64::from_datum(cst.constvalue, cst.constisnull).unwrap(); - return PgBox::::from_pg(ptr as _); - } - PgBox::::null() - }) - } - } -} - pub(crate) trait ReportableError { type Output; diff --git a/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs b/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs index 84b68fb6..e300f58c 100644 --- a/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs +++ b/wrappers/src/fdw/bigquery_fdw/bigquery_fdw.rs @@ -360,17 +360,6 @@ impl ForeignDataWrapper for BigQueryFdw { Ok(ret) } - fn get_rel_size( - &mut self, - _quals: &[Qual], - _columns: &[Column], - _sorts: &[Sort], - _limit: &Option, - _options: &HashMap, - ) -> Result<(i64, i32), BigQueryFdwError> { - Ok((0, 0)) - } - fn begin_scan( &mut self, quals: &[Qual], @@ -557,7 +546,7 @@ impl ForeignDataWrapper for BigQueryFdw { Ok(()) } - fn supported_aggregates(&self) -> Vec { + fn supported_aggregates() -> Vec { vec![ AggregateKind::Count, AggregateKind::CountColumn, @@ -568,7 +557,7 @@ impl ForeignDataWrapper for BigQueryFdw { ] } - fn supports_group_by(&self) -> bool { + fn supports_group_by() -> bool { true } diff --git a/wrappers/src/fdw/clickhouse_fdw/clickhouse_fdw.rs b/wrappers/src/fdw/clickhouse_fdw/clickhouse_fdw.rs index 316fab80..c1d373a1 100644 --- a/wrappers/src/fdw/clickhouse_fdw/clickhouse_fdw.rs +++ b/wrappers/src/fdw/clickhouse_fdw/clickhouse_fdw.rs @@ -1094,7 +1094,7 @@ impl ForeignDataWrapper for ClickHouseFdw { Ok(()) } - fn supported_aggregates(&self) -> Vec { + fn supported_aggregates() -> Vec { vec![ AggregateKind::Count, AggregateKind::CountColumn, @@ -1105,7 +1105,7 @@ impl ForeignDataWrapper for ClickHouseFdw { ] } - fn supports_group_by(&self) -> bool { + fn supports_group_by() -> bool { true } diff --git a/wrappers/src/fdw/mssql_fdw/mssql_fdw.rs b/wrappers/src/fdw/mssql_fdw/mssql_fdw.rs index deb59504..3c09751e 100644 --- a/wrappers/src/fdw/mssql_fdw/mssql_fdw.rs +++ b/wrappers/src/fdw/mssql_fdw/mssql_fdw.rs @@ -396,7 +396,7 @@ impl ForeignDataWrapper for MssqlFdw { Ok(()) } - fn supported_aggregates(&self) -> Vec { + fn supported_aggregates() -> Vec { vec![ AggregateKind::Count, AggregateKind::CountColumn, @@ -407,7 +407,7 @@ impl ForeignDataWrapper for MssqlFdw { ] } - fn supports_group_by(&self) -> bool { + fn supports_group_by() -> bool { true } diff --git a/wrappers/src/fdw/mysql_fdw/mysql_fdw.rs b/wrappers/src/fdw/mysql_fdw/mysql_fdw.rs index 8de1ee58..67b7044c 100644 --- a/wrappers/src/fdw/mysql_fdw/mysql_fdw.rs +++ b/wrappers/src/fdw/mysql_fdw/mysql_fdw.rs @@ -563,7 +563,7 @@ impl ForeignDataWrapper for MysqlFdw { self.disconnect_pool() } - fn supported_aggregates(&self) -> Vec { + fn supported_aggregates() -> Vec { vec![ AggregateKind::Count, AggregateKind::CountColumn, @@ -574,7 +574,7 @@ impl ForeignDataWrapper for MysqlFdw { ] } - fn supports_group_by(&self) -> bool { + fn supports_group_by() -> bool { true } diff --git a/wrappers/src/supabase_wrappers_tests.rs b/wrappers/src/supabase_wrappers_tests.rs index f9cb7523..58972533 100644 --- a/wrappers/src/supabase_wrappers_tests.rs +++ b/wrappers/src/supabase_wrappers_tests.rs @@ -1,12 +1,15 @@ -//! Runtime tests for `supabase-wrappers` core types that need a live Postgres backend -//! (e.g. `Cell::into_datum()`/`from_datum()` round trips). These can't run as plain -//! `cargo test` in the `supabase-wrappers` crate itself since it isn't a pgrx extension, -//! so they run here instead, against the real Postgres backend `cargo pgrx test` spins up. +//! Runtime tests for `supabase-wrappers` core types and framework behavior that need a +//! live Postgres backend (e.g. `Cell::into_datum()`/`from_datum()` round trips, or the +//! scan callback lifecycle under a cached plan). These can't run as plain `cargo test` +//! in the `supabase-wrappers` crate itself since it isn't a pgrx extension, so they run +//! here instead, against the real Postgres backend `cargo pgrx test` spins up. #[cfg(any(test, feature = "pg_test"))] #[pgrx::pg_schema] mod tests { + use pgrx::pg_sys::panic::ErrorReport; use pgrx::prelude::*; + use std::collections::HashMap; use supabase_wrappers::prelude::*; use supabase_wrappers::qual::form_array_from_datum; @@ -250,4 +253,264 @@ mod tests { let result = unsafe { form_array_from_datum(datum, false, pg_sys::UUIDARRAYOID) }; assert!(result.is_none()); } + + // ========================================================================== + // Regression test: cached-plan re-execution must not crash the backend + // ========================================================================== + + // `get_rel_size` is a planning-time-only trait hook and takes no `self` (planning + // never constructs an FDW instance), so it can't use instance-local state to detect + // repeat calls. Use a static counter instead to assert it's never re-run once the + // plan is cached. + static PLANNING_CALLS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + // Likewise, `new()` should now only run once per `EXECUTE` of the cached plan (never + // during planning), so this counter should end up equal to the number of executions. + static NEW_CALLS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0); + + #[wrappers_fdw( + version = "0.1.0", + author = "Supabase", + website = "https://github.com/supabase/wrappers", + error_type = "CacheTestFdwError" + )] + struct CacheTestFdw { + iter_done: bool, + tgt_cols: Vec, + } + + enum CacheTestFdwError {} + + impl From for ErrorReport { + fn from(_value: CacheTestFdwError) -> Self { + ErrorReport::new(PgSqlErrorCode::ERRCODE_FDW_ERROR, "", "") + } + } + + impl ForeignDataWrapper for CacheTestFdw { + fn new(_server: ForeignServer) -> Result { + NEW_CALLS.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(Self { + iter_done: false, + tgt_cols: Vec::new(), + }) + } + + // This method is called during the planning phase, we use it to + // assert that plan is being cached by checking that this is only + // ever called once. + fn get_rel_size( + _quals: &[Qual], + _columns: &[Column], + _sorts: &[Sort], + _limit: &Option, + _options: &HashMap, + ) -> Result<(i64, i32), CacheTestFdwError> { + PLANNING_CALLS.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok((0, 0)) + } + + fn begin_scan( + &mut self, + _quals: &[Qual], + columns: &[Column], + _sorts: &[Sort], + _limit: &Option, + _options: &HashMap, + ) -> Result<(), CacheTestFdwError> { + self.iter_done = false; + self.tgt_cols = columns.to_vec(); + Ok(()) + } + + fn iter_scan(&mut self, row: &mut Row) -> Result, CacheTestFdwError> { + if self.iter_done { + return Ok(None); + } + self.iter_done = true; + for col in &self.tgt_cols { + if col.name == "id" { + row.push("id", Some(Cell::I64(1))); + } + } + Ok(Some(())) + } + + fn end_scan(&mut self) -> Result<(), CacheTestFdwError> { + Ok(()) + } + } + + #[pg_test] + fn cached_plan_repeated_execution_does_not_crash() { + Spi::connect_mut(|c| { + c.update( + r#"create foreign data wrapper cache_test_wrapper + handler cache_test_fdw_handler validator cache_test_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"create server cache_test_server foreign data wrapper cache_test_wrapper"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"create foreign table cache_test_table (id bigint) server cache_test_server"#, + None, + &[], + ) + .unwrap(); + + // Use a prepared statement to force plan caching + c.update( + "prepare cache_test_q as select count(*) from cache_test_table", + None, + &[], + ) + .unwrap(); + + // Run the cached plan multiple times + for _ in 0..3 { + let count = c + .select("execute cache_test_q", None, &[]) + .unwrap() + .first() + .get_one::() + .unwrap(); + assert_eq!(count, Some(1)); + } + + c.update("deallocate cache_test_q", None, &[]).unwrap(); + + assert_eq!( + PLANNING_CALLS.load(std::sync::atomic::Ordering::SeqCst), + 1, + "expected get_rel_size to run exactly once for the whole cached plan" + ); + assert_eq!( + NEW_CALLS.load(std::sync::atomic::Ordering::SeqCst), + 3, + "expected new() to run exactly once per execution, never during planning" + ); + }); + } + + // ========================================================================== + // Regression test: https://github.com/supabase/wrappers/issues/237 + // ========================================================================== + + // Same underlying cause as a cached plan. This test is there just for completion + + #[wrappers_fdw( + version = "0.1.0", + author = "Supabase", + website = "https://github.com/supabase/wrappers", + error_type = "PlpgsqlCacheTestFdwError" + )] + struct PlpgsqlCacheTestFdw { + rows: Vec, + row_idx: usize, + } + + enum PlpgsqlCacheTestFdwError {} + + impl From for ErrorReport { + fn from(_value: PlpgsqlCacheTestFdwError) -> Self { + ErrorReport::new(PgSqlErrorCode::ERRCODE_FDW_ERROR, "", "") + } + } + + impl ForeignDataWrapper for PlpgsqlCacheTestFdw { + fn new(_server: ForeignServer) -> Result { + Ok(Self { + rows: vec![1, 2], + row_idx: 0, + }) + } + + fn begin_scan( + &mut self, + _quals: &[Qual], + _columns: &[Column], + _sorts: &[Sort], + _limit: &Option, + _options: &HashMap, + ) -> Result<(), PlpgsqlCacheTestFdwError> { + self.row_idx = 0; + Ok(()) + } + + fn iter_scan(&mut self, row: &mut Row) -> Result, PlpgsqlCacheTestFdwError> { + if self.row_idx >= self.rows.len() { + return Ok(None); + } + row.push("id", Some(Cell::I64(self.rows[self.row_idx]))); + self.row_idx += 1; + Ok(Some(())) + } + + fn end_scan(&mut self) -> Result<(), PlpgsqlCacheTestFdwError> { + Ok(()) + } + } + + #[pg_test] + fn plpgsql_function_wrapping_foreign_table_returns_consistent_results_across_calls() { + Spi::connect_mut(|c| { + c.update( + r#"create foreign data wrapper plpgsql_cache_test_wrapper + handler plpgsql_cache_test_fdw_handler validator plpgsql_cache_test_fdw_validator"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"create server plpgsql_cache_test_server foreign data wrapper plpgsql_cache_test_wrapper"#, + None, + &[], + ) + .unwrap(); + c.update( + r#"create foreign table plpgsql_cache_test_table (id bigint) server plpgsql_cache_test_server"#, + None, + &[], + ) + .unwrap(); + + // Mirrors the issue's `get_products()` repro: a plpgsql function whose + // body selects from the foreign table, called multiple times. + c.update( + r#"create function plpgsql_cache_test_get_ids() + returns table (id bigint) + language plpgsql as + $$ + begin + return query select t.id from plpgsql_cache_test_table t; + end + $$"#, + None, + &[], + ) + .unwrap(); + + for call in 0..3 { + let ids = c + .select( + "select id from plpgsql_cache_test_get_ids() order by id", + None, + &[], + ) + .unwrap() + .filter_map(|r| r.get_by_name::("id").unwrap()) + .collect::>(); + assert_eq!( + ids, + vec![1, 2], + "call #{call} to the cached plpgsql function returned wrong/missing rows" + ); + } + }); + } }