pgrx/
iter.rs

1//LICENSE Portions Copyright 2019-2021 ZomboDB, LLC.
2//LICENSE
3//LICENSE Portions Copyright 2021-2023 Technology Concepts & Design, Inc.
4//LICENSE
5//LICENSE Portions Copyright 2023-2023 PgCentral Foundation, Inc. <contact@pgcentral.org>
6//LICENSE
7//LICENSE All rights reserved.
8//LICENSE
9//LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
10#![deny(unsafe_op_in_unsafe_fn)]
11use core::{iter, ptr};
12
13use crate::callconv::{BoxRet, CallCx, RetAbi};
14use crate::fcinfo::{pg_return_null, srf_is_first_call, srf_return_done, srf_return_next};
15use crate::ptr::PointerExt;
16use crate::{pg_sys, IntoDatum, IntoHeapTuple, PgMemoryContexts};
17use pgrx_sql_entity_graph::metadata::{
18    ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
19};
20
21/// Support for returning a `SETOF T` from an SQL function.
22///
23/// [`SetOfIterator`] is typically used as a return type on `#[pg_extern]`-style functions
24/// and indicates that the SQL function should return a `SETOF` over the generic argument `T`.
25///
26/// It is a lightweight wrapper around an iterator, which you provide during construction.  The
27/// iterator *can* borrow from its environment, following Rust's normal borrowing rules.  If no
28/// borrowing is necessary, the `'static` lifetime should be used.
29///
30/// # Examples
31///
32/// This example simply returns a set of integers in the range `1..=5`.
33///
34/// ```rust,no_run
35/// use pgrx::prelude::*;
36/// #[pg_extern]
37/// fn return_ints() -> SetOfIterator<'static, i32> {
38///     SetOfIterator::new(1..=5)
39/// }
40/// ```
41///
42/// Here we return a set of `&str`s, borrowed from an argument:
43///
44/// ```rust,no_run
45/// use pgrx::prelude::*;
46/// #[pg_extern]
47/// fn split_string<'a>(input: &'a str) -> SetOfIterator<'a, &'a str> {
48///     SetOfIterator::new(input.split_whitespace())
49/// }
50/// ```
51#[repr(transparent)]
52pub struct SetOfIterator<'a, T>(
53    // Postgres uses the same ABI for `returns setof` and 1-column `returns table`
54    TableIterator<'a, (T,)>,
55);
56
57impl<'a, T: 'a> SetOfIterator<'a, T> {
58    pub fn new(iter: impl IntoIterator<Item = T> + 'a) -> Self {
59        // Forward the impl using remapping to minimize `unsafe` code and keep them in sync
60        Self(TableIterator::new(iter.into_iter().map(|c| (c,))))
61    }
62
63    pub fn empty() -> Self {
64        Self::new(iter::empty())
65    }
66
67    pub fn once(value: T) -> Self {
68        Self::new(iter::once(value))
69    }
70}
71
72impl<'a, T> Iterator for SetOfIterator<'a, T> {
73    type Item = T;
74
75    #[inline]
76    fn next(&mut self) -> Option<Self::Item> {
77        self.0.next().map(|(val,)| val)
78    }
79}
80
81/// `SetOfIterator<'_, T>` differs from `TableIterator<'a, (T,)>` in generated SQL
82unsafe impl<'a, T> SqlTranslatable for SetOfIterator<'a, T>
83where
84    T: SqlTranslatable,
85{
86    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
87        T::argument_sql()
88    }
89    fn return_sql() -> Result<Returns, ReturnsError> {
90        match T::return_sql() {
91            Ok(Returns::One(sql)) => Ok(Returns::SetOf(sql)),
92            Ok(Returns::SetOf(_)) => Err(ReturnsError::NestedSetOf),
93            Ok(Returns::Table(_)) => Err(ReturnsError::SetOfContainingTable),
94            err @ Err(_) => err,
95        }
96    }
97}
98
99/// Support for a `TABLE (...)` from an SQL function.
100///
101/// [`TableIterator`] is typically used as the return type of a `#[pg_extern]`-style function,
102/// indicating that the function returns a table of named columns.  [`TableIterator`] is
103/// generic over `Row`, but that `Row` must be a Rust tuple containing one or more elements.  They
104/// must also be "named" using pgrx's [`name!`][crate::name] macro.  See the examples below.
105///
106/// It is a lightweight wrapper around an iterator, which you provide during construction.  The
107/// iterator *can* borrow from its environment, following Rust's normal borrowing rules.  If no
108/// borrowing is necessary, the `'static` lifetime should be used.
109///
110/// # Examples
111///
112/// This example returns a table of employee information.
113///
114/// ```rust,no_run
115/// use pgrx::prelude::*;
116/// #[pg_extern]
117/// fn employees() -> TableIterator<'static,
118///         (
119///             name!(id, i64),
120///             name!(dept_code, String),
121///             name!(full_text, &'static str)
122///         )
123/// > {
124///     TableIterator::new(vec![
125///         (42, "ARQ".into(), "John Hammond"),
126///         (87, "EGA".into(), "Mary Kinson"),
127///         (3,  "BLA".into(), "Perry Johnson"),
128///     ])
129/// }
130/// ```
131///
132/// And here we return a simple numbered list of words, borrowed from the input `&str`.
133///
134/// ```rust,no_run
135/// use pgrx::prelude::*;
136/// #[pg_extern]
137/// fn split_string<'a>(input: &'a str) -> TableIterator<'a, ( name!(num, i32), name!(word, &'a str) )> {
138///     TableIterator::new(input.split_whitespace().enumerate().map(|(n, w)| (n as i32, w)))
139/// }
140/// ```
141pub struct TableIterator<'a, Row> {
142    iter: Box<dyn Iterator<Item = Row> + 'a>,
143}
144
145impl<'a, Row: 'a> TableIterator<'a, Row> {
146    pub fn new(iter: impl IntoIterator<Item = Row> + 'a) -> Self {
147        Self { iter: Box::new(iter.into_iter()) }
148    }
149
150    pub fn empty() -> Self {
151        Self::new(iter::empty())
152    }
153
154    pub fn once(value: Row) -> Self {
155        Self::new(iter::once(value))
156    }
157}
158
159impl<'a, Row> Iterator for TableIterator<'a, Row> {
160    type Item = Row;
161
162    #[inline]
163    fn next(&mut self) -> Option<Self::Item> {
164        self.iter.next()
165    }
166}
167
168unsafe impl<'iter, C> SqlTranslatable for TableIterator<'iter, (C,)>
169where
170    C: SqlTranslatable + 'iter,
171{
172    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
173        Err(ArgumentError::Table)
174    }
175    fn return_sql() -> Result<Returns, ReturnsError> {
176        let vec = vec![C::return_sql().and_then(|sql| match sql {
177            Returns::One(sql) => Ok(sql),
178            Returns::SetOf(_) => Err(ReturnsError::TableContainingSetOf),
179            Returns::Table(_) => Err(ReturnsError::NestedTable),
180        })?];
181        Ok(Returns::Table(vec))
182    }
183}
184
185unsafe impl<'a, T> RetAbi for SetOfIterator<'a, T>
186where
187    T: BoxRet,
188{
189    type Item = <Self as Iterator>::Item;
190    type Ret = IterRet<Self>;
191
192    unsafe fn check_fcinfo_and_prepare(fcinfo: pg_sys::FunctionCallInfo) -> CallCx {
193        unsafe { TableIterator::<(T,)>::check_fcinfo_and_prepare(fcinfo) }
194    }
195
196    fn to_ret(self) -> Self::Ret {
197        let mut iter = self;
198        IterRet(match iter.next() {
199            None => Step::Done,
200            Some(value) => Step::Init(iter, value),
201        })
202    }
203
204    unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum {
205        let ret = match ret.0 {
206            Step::Done => Step::Done,
207            Step::Once(value) => Step::Once((value,)),
208            Step::Init(iter, value) => Step::Init(iter.0, (value,)),
209        };
210        unsafe { TableIterator::<(T,)>::box_ret_in_fcinfo(fcinfo, IterRet(ret)) }
211    }
212
213    unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo) {}
214
215    unsafe fn move_into_fcinfo_fcx(self, fcinfo: pg_sys::FunctionCallInfo) {
216        unsafe { self.0.move_into_fcinfo_fcx(fcinfo) }
217    }
218
219    unsafe fn ret_from_fcinfo_fcx(fcinfo: pg_sys::FunctionCallInfo) -> Self::Ret {
220        let step = match unsafe { TableIterator::<(T,)>::ret_from_fcinfo_fcx(fcinfo).0 } {
221            Step::Done => Step::Done,
222            Step::Once((item,)) => Step::Once(item),
223            Step::Init(iter, (value,)) => Step::Init(Self(iter), value),
224        };
225        IterRet(step)
226    }
227
228    unsafe fn finish_call_fcinfo(fcinfo: pg_sys::FunctionCallInfo) {
229        unsafe { TableIterator::<(T,)>::finish_call_fcinfo(fcinfo) }
230    }
231}
232
233unsafe impl<'a, Row> RetAbi for TableIterator<'a, Row>
234where
235    Row: RetAbi,
236{
237    type Item = <Self as Iterator>::Item;
238    type Ret = IterRet<Self>;
239
240    unsafe fn check_fcinfo_and_prepare(fcinfo: pg_sys::FunctionCallInfo) -> CallCx {
241        unsafe {
242            if srf_is_first_call(fcinfo) {
243                let fn_call_cx = pg_sys::init_MultiFuncCall(fcinfo);
244                CallCx::WrappedFn((*fn_call_cx).multi_call_memory_ctx)
245            } else {
246                CallCx::RestoreCx
247            }
248        }
249    }
250
251    fn to_ret(self) -> Self::Ret {
252        let mut iter = self;
253        IterRet(match iter.next() {
254            None => Step::Done,
255            Some(value) => Step::Init(iter, value),
256        })
257    }
258
259    unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum {
260        let value = unsafe {
261            match ret.0 {
262                Step::Done => return empty_srf(fcinfo),
263                Step::Once(value) => value,
264                Step::Init(iter, value) => {
265                    // Move the iterator in and don't worry about putting it back
266                    iter.move_into_fcinfo_fcx(fcinfo);
267                    value.fill_fcinfo_fcx(fcinfo);
268                    value
269                }
270            }
271        };
272
273        unsafe {
274            let fcx = deref_fcx(fcinfo);
275            srf_return_next(fcinfo, fcx);
276            <Row as RetAbi>::box_ret_in_fcinfo(fcinfo, value.to_ret())
277        }
278    }
279
280    unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo) {}
281
282    unsafe fn move_into_fcinfo_fcx(self, fcinfo: pg_sys::FunctionCallInfo) {
283        unsafe {
284            let fcx = deref_fcx(fcinfo);
285            let ptr = srf_memcx(fcx).leak_and_drop_on_delete(self);
286            // it's the first call so we need to finish setting up fcx
287            (*fcx).user_fctx = ptr.cast();
288        }
289    }
290
291    unsafe fn ret_from_fcinfo_fcx(fcinfo: pg_sys::FunctionCallInfo) -> Self::Ret {
292        // SAFETY: fcx.user_fctx was set earlier, immediately before or in a prior call
293        let iter = unsafe {
294            let fcx = deref_fcx(fcinfo);
295            &mut *(*fcx).user_fctx.cast::<TableIterator<Row>>()
296        };
297        IterRet(match iter.next() {
298            None => Step::Done,
299            Some(value) => Step::Once(value),
300        })
301    }
302
303    unsafe fn finish_call_fcinfo(fcinfo: pg_sys::FunctionCallInfo) {
304        unsafe {
305            let fcx = deref_fcx(fcinfo);
306            srf_return_done(fcinfo, fcx)
307        }
308    }
309}
310
311/// How iterators are returned
312pub struct IterRet<T: RetAbi>(Step<T>);
313
314/// ValuePerCall SRF steps
315enum Step<T: RetAbi> {
316    Done,
317    Once(T::Item),
318    Init(T, T::Item),
319}
320
321pub(crate) unsafe fn empty_srf(fcinfo: pg_sys::FunctionCallInfo) -> pg_sys::Datum {
322    unsafe {
323        let fcx = deref_fcx(fcinfo);
324        srf_return_done(fcinfo, fcx);
325        pg_return_null(fcinfo)
326    }
327}
328
329/// "per_MultiFuncCall" but no FFI cost
330pub(crate) unsafe fn deref_fcx(fcinfo: pg_sys::FunctionCallInfo) -> *mut pg_sys::FuncCallContext {
331    unsafe { (*(*fcinfo).flinfo).fn_extra.cast() }
332}
333
334pub(crate) unsafe fn srf_memcx(fcx: *mut pg_sys::FuncCallContext) -> PgMemoryContexts {
335    unsafe { PgMemoryContexts::For((*fcx).multi_call_memory_ctx) }
336}
337
338/// Return ABI for single-column tuples
339///
340/// Postgres extended SQL with `returns setof $ty` before SQL had `returns table($($ident $ty),*)`.
341/// Due to this history, single-column `returns table` reuses the internals of `returns setof`.
342/// This lets them simply return a simple Datum instead of handling a TupleDesc and HeapTuple, but
343/// means we need to have this distinct impl, as the return type is not `TYPEFUNC_COMPOSITE`!
344/// Fortunately, RetAbi lets `TableIterator<'a, Tup>` handle this by calling `<Tup as RetAbi>`.
345unsafe impl<C> RetAbi for (C,)
346where
347    C: BoxRet, // so we support TableIterator<'a, (Option<T>,)> as well
348{
349    type Item = C;
350    type Ret = C;
351
352    fn to_ret(self) -> Self::Ret {
353        self.0
354    }
355
356    /// Returning a "row" of only one column is identical to returning that single type
357    unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum {
358        unsafe { C::box_ret_in_fcinfo(fcinfo, ret.to_ret()) }
359    }
360
361    unsafe fn fill_fcinfo_fcx(&self, _fcinfo: pg_sys::FunctionCallInfo) {}
362
363    unsafe fn move_into_fcinfo_fcx(self, _fcinfo: pg_sys::FunctionCallInfo) {}
364}
365
366macro_rules! impl_table_iter {
367    ($($C:ident),* $(,)?) => {
368        unsafe impl<'iter, $($C,)*> SqlTranslatable for TableIterator<'iter, ($($C,)*)>
369        where
370            $($C: SqlTranslatable + 'iter,)*
371        {
372            fn argument_sql() -> Result<SqlMapping, ArgumentError> {
373                Err(ArgumentError::Table)
374            }
375            fn return_sql() -> Result<Returns, ReturnsError> {
376                let vec = vec![
377                $(
378                    $C::return_sql().and_then(|sql| match sql {
379                        Returns::One(sql) => Ok(sql),
380                        Returns::SetOf(_) => Err(ReturnsError::TableContainingSetOf),
381                        Returns::Table(_) => Err(ReturnsError::NestedTable),
382                    })?,
383                )*
384                ];
385                Ok(Returns::Table(vec))
386            }
387        }
388
389        impl<$($C: IntoDatum),*> IntoHeapTuple for ($($C,)*) {
390            unsafe fn into_heap_tuple(self, tupdesc: pg_sys::TupleDesc) -> *mut pg_sys::HeapTupleData {
391                // shadowing the type names with these identifiers
392                #[allow(nonstandard_style)]
393                let ($($C,)*) = self;
394                let datums = [$($C.into_datum(),)*];
395                let mut nulls = datums.map(|option| option.is_none());
396                let mut datums = datums.map(|option| option.unwrap_or(pg_sys::Datum::from(0)));
397
398                unsafe {
399                    // SAFETY:  Caller has asserted that `tupdesc` is valid, and we just went
400                    // through a little bit of effort to setup properly sized arrays for
401                    // `datums` and `nulls`
402                    pg_sys::heap_form_tuple(tupdesc, datums.as_mut_ptr(), nulls.as_mut_ptr())
403                }
404            }
405        }
406
407        unsafe impl<$($C),*> RetAbi for ($($C,)*)
408        where
409             $($C: BoxRet,)*
410             Self: IntoHeapTuple,
411        {
412            type Item = Self;
413            type Ret = Self;
414
415            fn to_ret(self) -> Self::Ret {
416                self
417            }
418
419            unsafe fn box_ret_in_fcinfo(fcinfo: pg_sys::FunctionCallInfo, ret: Self::Ret) -> pg_sys::Datum {
420                unsafe {
421                    let fcx = deref_fcx(fcinfo);
422                    let heap_tuple = ret.into_heap_tuple((*fcx).tuple_desc);
423                    pg_sys::HeapTupleHeaderGetDatum((*heap_tuple).t_data)
424                }
425            }
426
427            unsafe fn move_into_fcinfo_fcx(self, _fcinfo: pg_sys::FunctionCallInfo) {}
428
429            unsafe fn fill_fcinfo_fcx(&self, fcinfo: pg_sys::FunctionCallInfo) {
430                // Pure side effect, leave the value in place.
431                unsafe {
432                    let fcx = deref_fcx(fcinfo);
433                    srf_memcx(fcx).switch_to(|_| {
434                        let mut tupdesc = ptr::null_mut();
435                        let mut oid = pg_sys::Oid::default();
436                        let ty_class = pg_sys::get_call_result_type(fcinfo, &mut oid, &mut tupdesc);
437                        if tupdesc.is_non_null() && ty_class == pg_sys::TypeFuncClass::TYPEFUNC_COMPOSITE {
438                            pg_sys::BlessTupleDesc(tupdesc);
439                            (*fcx).tuple_desc = tupdesc;
440                        }
441                    });
442                }
443            }
444        }
445
446    }
447}
448
449impl_table_iter!(T0, T1);
450impl_table_iter!(T0, T1, T2);
451impl_table_iter!(T0, T1, T2, T3);
452impl_table_iter!(T0, T1, T2, T3, T4);
453impl_table_iter!(T0, T1, T2, T3, T4, T5);
454impl_table_iter!(T0, T1, T2, T3, T4, T5, T6);
455impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7);
456impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8);
457impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9);
458impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10);
459impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11);
460impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12);
461impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13);
462impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14);
463impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15);
464impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16);
465impl_table_iter!(T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17);
466impl_table_iter!(
467    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18
468);
469impl_table_iter!(
470    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19
471);
472impl_table_iter!(
473    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20
474);
475impl_table_iter!(
476    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
477    T21
478);
479impl_table_iter!(
480    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
481    T21, T22
482);
483impl_table_iter!(
484    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
485    T21, T22, T23
486);
487impl_table_iter!(
488    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
489    T21, T22, T23, T24
490);
491impl_table_iter!(
492    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
493    T21, T22, T23, T24, T25
494);
495impl_table_iter!(
496    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
497    T21, T22, T23, T24, T25, T26
498);
499impl_table_iter!(
500    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
501    T21, T22, T23, T24, T25, T26, T27
502);
503impl_table_iter!(
504    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
505    T21, T22, T23, T24, T25, T26, T27, T28
506);
507impl_table_iter!(
508    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
509    T21, T22, T23, T24, T25, T26, T27, T28, T29
510);
511impl_table_iter!(
512    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
513    T21, T22, T23, T24, T25, T26, T27, T28, T29, T30
514);
515impl_table_iter!(
516    T0, T1, T2, T3, T4, T5, T6, T7, T8, T9, T10, T11, T12, T13, T14, T15, T16, T17, T18, T19, T20,
517    T21, T22, T23, T24, T25, T26, T27, T28, T29, T30, T31
518);
pgrx/iter.rs

pgrx/
iter.rs
