Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions CLAUDE.md
Original file line number Diff line number Diff line change
Expand Up @@ -53,12 +53,13 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
fn end_modify(&mut self) -> Result<(), E>;

// Optional methods for aggregate pushdown
fn supported_aggregates(&self) -> Vec<AggregateKind>;
fn supports_group_by(&self) -> bool;
fn supported_aggregates() -> Vec<AggregateKind>;
fn supports_group_by() -> bool;
fn begin_aggregate_scan(&mut self, aggregates: &[Aggregate], group_by: &[Column], quals: &[Qual], options: &HashMap<String, String>) -> 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<Vec<String>, E>;
fn validator(options: Vec<Option<String>>, catalog: Option<Oid>) -> Result<(), E>;
Expand Down Expand Up @@ -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<AggregateKind> {
fn supported_aggregates() -> Vec<AggregateKind> {
vec![
AggregateKind::Count,
AggregateKind::CountColumn,
Expand All @@ -329,7 +330,7 @@ fn supported_aggregates(&self) -> Vec<AggregateKind> {
]
}

fn supports_group_by(&self) -> bool { true }
fn supports_group_by() -> bool { true }

fn begin_aggregate_scan(
&mut self,
Expand Down
4 changes: 2 additions & 2 deletions docs/guides/query-pushdown.md
Original file line number Diff line number Diff line change
Expand Up @@ -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<AggregateKind> {
fn supported_aggregates() -> Vec<AggregateKind> {
vec![AggregateKind::Count, AggregateKind::Sum, AggregateKind::Avg]
}

fn supports_group_by(&self) -> bool {
fn supports_group_by() -> bool {
true
}

Expand Down
40 changes: 34 additions & 6 deletions supabase-wrappers/src/interface.rs
Original file line number Diff line number Diff line change
Expand Up @@ -551,6 +551,15 @@ pub struct Qual {
pub value: Value,
pub use_or: bool,
pub param: Option<Param>,

// 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<usize>,
}

impl Qual {
Expand Down Expand Up @@ -881,6 +890,12 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
/// 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<Self, E>
where
Self: Sized;
Expand All @@ -890,9 +905,12 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
/// 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],
Expand Down Expand Up @@ -1018,10 +1036,13 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
///
/// ## 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<AggregateKind> {
/// fn supported_aggregates() -> Vec<AggregateKind> {
/// vec![
/// AggregateKind::Count,
/// AggregateKind::CountColumn,
Expand All @@ -1032,7 +1053,7 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
/// ]
/// }
/// ```
fn supported_aggregates(&self) -> Vec<AggregateKind> {
fn supported_aggregates() -> Vec<AggregateKind> {
vec![]
}

Expand All @@ -1043,14 +1064,17 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
///
/// 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
}

Expand Down Expand Up @@ -1184,7 +1208,11 @@ pub trait ForeignDataWrapper<E: Into<ErrorReport>> {
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
Expand Down
31 changes: 14 additions & 17 deletions supabase-wrappers/src/qual.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -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);
}
Expand Down Expand Up @@ -296,14 +297,14 @@ 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)
}
}

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,
Expand Down Expand Up @@ -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);
}
Expand All @@ -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,
Expand All @@ -385,14 +386,14 @@ pub(crate) unsafe fn extract_from_var(
value: Value::Cell(Cell::Bool(true)),
use_or: false,
param: None,
value_const: None,
};

Some(qual)
}
}

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,
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -456,14 +458,14 @@ pub(crate) unsafe fn extract_from_boolean_test(
value: Value::Cell(Cell::Bool(value)),
use_or: false,
param: None,
value_const: None,
};

Some(qual)
}
}

pub(crate) unsafe fn extract_quals(
root: *mut pg_sys::PlannerInfo,
baserel: *mut pg_sys::RelOptInfo,
baserel_id: pg_sys::Oid,
) -> Vec<Qual> {
Expand All @@ -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 {
Expand Down
Loading