1//! Refuses a plan before it runs if any call to the extension's functions
2//! was left outside a jev scan. Such a call would run once per row, and
3//! its body refuses; but by then the plan's scans may already have judged
4//! a window of rows, and that is money spent on a failed statement. The
5//! check runs in `ExecutorStart`, after the plan is initialized and before
6//! any row is pulled, so nothing has been sent.
7//!
8//! The same check then prices the plan: the rows every jev scan's child
9//! is estimated to produce go to [`Judge::afford`], which refuses a
10//! statement over its spend guards before its first request.
11
12use std::ffi::c_void;
13use std::sync::OnceLock;
14
15use pgrx::pg_sys::panic::ErrorReport;
16use pgrx::{PgLogLevel, PgSqlErrorCode, pg_guard, pg_sys};
17
18use super::ffi::ptrs;
19use super::{Estimate, Judge};
20use super::plan::{contains_call, function_oids, tables};
21
22static PREVIOUS: OnceLock<pg_sys::ExecutorStart_hook_type> = OnceLock::new();
23
24/// # Safety
25/// Called once, from `_PG_init`.
26pub(super) unsafe fn install<J: Judge>() {
27    unsafe {
28        let _ = PREVIOUS.set(pg_sys::ExecutorStart_hook);
29        pg_sys::ExecutorStart_hook = Some(executor_start::<J>);
30    }
31}
32
33#[pg_guard]
34unsafe extern "C-unwind" fn executor_start<J: Judge>(query: *mut pg_sys::QueryDesc, eflags: i32) {
35    // SAFETY: the executor's arguments; the plan state exists once the
36    // standard start returns.
37    unsafe {
38        match PREVIOUS.get() {
39            Some(Some(previous)) => previous(query, eflags),
40            _ => pg_sys::standard_ExecutorStart(query, eflags),
41        }
42        // EXPLAIN without ANALYZE spends nothing, and showing the plan is
43        // how a user finds the call that would be refused.
44        if eflags as u32 & pg_sys::EXEC_FLAG_EXPLAIN_ONLY != 0 {
45            return;
46        }
47        let Some(oids) = function_oids::<J>() else { return };
48        if let Some(tag) = unbatched((*query).planstate, &oids) {
49            let tag = format!("{tag:?}");
50            ErrorReport::new(
51                PgSqlErrorCode::ERRCODE_FEATURE_NOT_SUPPORTED,
52                format!(
53                    "a jev call in this query is outside the jev scan, in a {} node",
54                    tag.trim_start_matches("T_")
55                ),
56                "jev scan",
57            )
58            .set_detail("It would send one request per row, one at a time. Nothing was sent.")
59            .set_hint("Give each call the columns of one relation, outside a join's select list and conditions.")
60            .report(PgLogLevel::ERROR);
61        }
62        if let Err(refusal) = J::afford(&estimate((*query).planstate)) {
63            refusal.report(PgLogLevel::ERROR);
64        }
65    }
66}
67
68/// What the plan's jev scans are estimated to judge: their children's
69/// rows, which is after every SQL filter. A scan rescanned under a nested
70/// loop counts once; the scan's own count while running backs that up.
71unsafe fn estimate(root: *mut pg_sys::PlanState) -> Estimate {
72    #[pg_guard]
73    unsafe extern "C-unwind" fn visit(state: *mut pg_sys::PlanState, context: *mut c_void) -> bool {
74        // SAFETY: `context` is the Estimate below, alive for the walk; a
75        // jev scan's plan has its one child in `custom_plans`.
76        unsafe {
77            let estimate = &mut *context.cast::<Estimate>();
78            let plan = (*state).plan;
79            if (*plan).type_ == pg_sys::NodeTag::T_CustomScan {
80                let scan = plan.cast::<pg_sys::CustomScan>();
81                if std::ptr::eq((*scan).methods, &tables().scan) {
82                    let one = scan_estimate(scan);
83                    estimate.rows += one.rows;
84                    estimate.bytes += one.bytes;
85                }
86            }
87            pg_sys::planstate_tree_walker_impl(state, Some(visit), context)
88        }
89    }
90
91    let mut estimate = Estimate::default();
92    // SAFETY: an initialized plan state tree; `visit` never raises.
93    unsafe { visit(root, (&raw mut estimate).cast()) };
94    estimate
95}
96
97/// What one jev scan is estimated to judge: its child's rows and width.
98///
99/// # Safety
100/// `scan` is a jev scan's plan, with its one child in `custom_plans`.
101pub(super) unsafe fn scan_estimate(scan: *mut pg_sys::CustomScan) -> Estimate {
102    unsafe {
103        let child = ptrs::<pg_sys::Plan>((*scan).custom_plans).next().expect("one child plan");
104        Estimate { rows: (*child).plan_rows, bytes: (*child).plan_rows * f64::from((*child).plan_width) }
105    }
106}
107
108/// The tag of the first plan node that evaluates a call itself.
109unsafe fn unbatched(root: *mut pg_sys::PlanState, oids: &[pg_sys::Oid]) -> Option<pg_sys::NodeTag> {
110    struct Search<'a> {
111        oids: &'a [pg_sys::Oid],
112        found: Option<pg_sys::NodeTag>,
113    }
114
115    #[pg_guard]
116    unsafe extern "C-unwind" fn visit(state: *mut pg_sys::PlanState, context: *mut c_void) -> bool {
117        // SAFETY: `context` is the Search below, alive for the walk; the
118        // plan states and their plans belong to the running executor.
119        unsafe {
120            let search = &mut *context.cast::<Search>();
121            let plan = (*state).plan;
122            if expressions(plan).into_iter().any(|e| contains_call(e, search.oids)) {
123                search.found = Some((*plan).type_);
124                return true;
125            }
126            pg_sys::planstate_tree_walker_impl(state, Some(visit), context)
127        }
128    }
129
130    let mut search = Search { oids, found: None };
131    // SAFETY: an initialized plan state tree; `visit` never raises.
132    unsafe { visit(root, (&raw mut search).cast()) };
133    search.found
134}
135
136/// Where a plan node keeps expressions it evaluates per row. A jev scan's
137/// own call columns are descriptors, never evaluated, so only its call
138/// arguments count.
139unsafe fn expressions(plan: *mut pg_sys::Plan) -> Vec<*mut pg_sys::Node> {
140    use pg_sys::NodeTag::*;
141    unsafe {
142        let mut found: Vec<*mut pg_sys::Node> = vec![(*plan).targetlist.cast(), (*plan).qual.cast()];
143        match (*plan).type_ {
144            T_NestLoop | T_MergeJoin | T_HashJoin => found.push((*plan.cast::<pg_sys::Join>()).joinqual.cast()),
145            T_Result => found.push((*plan.cast::<pg_sys::Result>()).resconstantqual),
146            T_FunctionScan => found.push((*plan.cast::<pg_sys::FunctionScan>()).functions.cast()),
147            T_ValuesScan => found.push((*plan.cast::<pg_sys::ValuesScan>()).values_lists.cast()),
148            // RETURNING and ON CONFLICT are evaluated per row, above the
149            // scan; MERGE's actions likewise.
150            T_ModifyTable => {
151                let modify = plan.cast::<pg_sys::ModifyTable>();
152                found.push((*modify).returningLists.cast());
153                found.push((*modify).onConflictSet.cast());
154                found.push((*modify).onConflictWhere);
155                found.push((*modify).mergeActionLists.cast());
156            }
157            T_CustomScan => {
158                let scan = plan.cast::<pg_sys::CustomScan>();
159                if std::ptr::eq((*scan).methods, &tables().scan) {
160                    found.push((*scan).custom_exprs.cast());
161                }
162            }
163            _ => {}
164        }
165        // NULL is an empty list; a FunctionScan's RangeTblFunctions are
166        // nodes the expression walker descends through to their funcexpr.
167        found.retain(|e| !e.is_null());
168        found
169    }
170}