pgx/
iter.rs

1use std::iter::once;
2
3use crate::IntoHeapTuple;
4use pgx_sql_entity_graph::metadata::{
5    ArgumentError, Returns, ReturnsError, SqlMapping, SqlTranslatable,
6};
7
8/// Support for returning a `SETOF T` from an SQL function.
9///
10/// [`SetOfIterator`] is typically used as a return type on `#[pg_extern]`-style functions
11/// and indicates that the SQL function should return a `SETOF` over the generic argument `T`.
12///
13/// It is a lightweight wrapper around an iterator, which you provide during construction.  The
14/// iterator *can* borrow from its environment, following Rust's normal borrowing rules.  If no
15/// borrowing is necessary, the `'static` lifetime should be used.
16///
17/// # Examples
18///
19/// This example simply returns a set of integers in the range `1..=5`.
20///
21/// ```rust,no_run
22/// use pgx::prelude::*;
23/// #[pg_extern]
24/// fn return_ints() -> SetOfIterator<'static, i32> {
25///     SetOfIterator::new(1..=5)
26/// }
27/// ```
28///
29/// Here we return a set of `&str`s, borrowed from an argument:
30///
31/// ```rust,no_run
32/// use pgx::prelude::*;
33/// #[pg_extern]
34/// fn split_string<'a>(input: &'a str) -> SetOfIterator<'a, &'a str> {
35///     SetOfIterator::new(input.split_whitespace())
36/// }
37/// ```
38pub struct SetOfIterator<'a, T> {
39    iter: Box<dyn Iterator<Item = T> + 'a>,
40}
41
42impl<'a, T> SetOfIterator<'a, T> {
43    pub fn new<I>(iter: I) -> Self
44    where
45        I: IntoIterator<Item = T> + 'a,
46    {
47        Self { iter: Box::new(iter.into_iter()) }
48    }
49}
50
51impl<'a, T> Iterator for SetOfIterator<'a, T> {
52    type Item = T;
53
54    #[inline]
55    fn next(&mut self) -> Option<Self::Item> {
56        self.iter.next()
57    }
58}
59
60unsafe impl<'a, T> SqlTranslatable for SetOfIterator<'a, T>
61where
62    T: SqlTranslatable,
63{
64    fn argument_sql() -> Result<SqlMapping, ArgumentError> {
65        T::argument_sql()
66    }
67    fn return_sql() -> Result<Returns, ReturnsError> {
68        match T::return_sql() {
69            Ok(Returns::One(sql)) => Ok(Returns::SetOf(sql)),
70            Ok(Returns::SetOf(_)) => Err(ReturnsError::NestedSetOf),
71            Ok(Returns::Table(_)) => Err(ReturnsError::SetOfContainingTable),
72            err @ Err(_) => err,
73        }
74    }
75}
76
77/// Support for a `TABLE (...)` from an SQL function.
78///
79/// [`TableIterator`] is typically used as the return type of a `#[pg_extern]`-style function,
80/// indicating that the function returns a table of named columns.  [`TableIterator`] is
81/// generic over `T`, but that `T` must be a Rust tuple containing one or more elements.  They
82/// must also be "named" using pgx' [`name!`] macro.  See the examples below.
83///
84/// It is a lightweight wrapper around an iterator, which you provide during construction.  The
85/// iterator *can* borrow from its environment, following Rust's normal borrowing rules.  If no
86/// borrowing is necessary, the `'static` lifetime should be used.
87///
88/// # Examples
89///
90/// This example returns a table of employee information.
91///
92/// ```rust,no_run
93/// use pgx::prelude::*;
94/// #[pg_extern]
95/// fn employees() -> TableIterator<'static,
96///         (
97///             name!(id, i64),
98///             name!(dept_code, String),
99///             name!(full_text, &'static str)
100///         )
101/// > {
102///     TableIterator::new(vec![
103///         (42, "ARQ".into(), "John Hammond"),
104///         (87, "EGA".into(), "Mary Kinson"),
105///         (3,  "BLA".into(), "Perry Johnson"),
106///     ])
107/// }
108/// ```
109///
110/// And here we return a simple numbered list of words, borrowed from the input `&str`.
111///
112/// ```rust,no_run
113/// use pgx::prelude::*;
114/// #[pg_extern]
115/// fn split_string<'a>(input: &'a str) -> TableIterator<'a, ( name!(num, i32), name!(word, &'a str) )> {
116///     TableIterator::new(input.split_whitespace().enumerate().map(|(n, w)| (n as i32, w)))
117/// }
118/// ```
119pub struct TableIterator<'a, T> {
120    iter: Box<dyn Iterator<Item = T> + 'a>,
121}
122
123impl<'a, T> TableIterator<'a, T>
124where
125    T: IntoHeapTuple + 'a,
126{
127    pub fn new<I>(iter: I) -> Self
128    where
129        I: IntoIterator<Item = T> + 'a,
130    {
131        Self { iter: Box::new(iter.into_iter()) }
132    }
133
134    pub fn once(value: T) -> Self {
135        Self::new(once(value))
136    }
137}
138
139impl<'a, T> Iterator for TableIterator<'a, T> {
140    type Item = T;
141
142    #[inline]
143    fn next(&mut self) -> Option<Self::Item> {
144        self.iter.next()
145    }
146}
147
148seq_macro::seq!(I in 0..=32 {
149    #(
150        seq_macro::seq!(N in 0..=I {
151            unsafe impl<'a, #(Input~N,)*> SqlTranslatable for TableIterator<'a, (#(Input~N,)*)>
152            where
153                #(
154                    Input~N: SqlTranslatable + 'a,
155                )*
156            {
157                fn argument_sql() -> Result<SqlMapping, ArgumentError> {
158                    Err(ArgumentError::Table)
159                }
160                fn return_sql() -> Result<Returns, ReturnsError> {
161                    let mut vec = Vec::new();
162                    #(
163                        vec.push(match Input~N::return_sql() {
164                            Ok(Returns::One(sql)) => sql,
165                            Ok(Returns::SetOf(_)) => return Err(ReturnsError::TableContainingSetOf),
166                            Ok(Returns::Table(_)) => return Err(ReturnsError::NestedTable),
167                            Err(err) => return Err(err),
168                        });
169                    )*
170                    Ok(Returns::Table(vec))
171                }
172            }
173        });
174    )*
175});
pgx/iter.rs

pgx/
iter.rs
