Skip to content

Commit 01d1d24

Browse files
committed
Problem: cursors may or may not be readonly-friendly
However, we always assume them to be non-readonly. Solution: split the API to `open_cursor` and `open_cursor_mut` This way we can ensure we're squeezing out performance in the right place, when feasible.
1 parent 49f1037 commit 01d1d24

File tree

2 files changed

+44
-1
lines changed

2 files changed

+44
-1
lines changed

pgx-tests/src/tests/spi_tests.rs

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,27 @@ mod tests {
209209
});
210210
}
211211

212+
#[pg_test]
213+
fn test_cursor_mut() {
214+
Spi::execute(|mut client| {
215+
client.update("CREATE TABLE tests.cursor_table (id int)", None, None);
216+
217+
let mut portal = client.open_cursor_mut(
218+
"INSERT INTO tests.cursor_table (id) \
219+
SELECT i FROM generate_series(1, 10) AS t(i) RETURNING id",
220+
None,
221+
);
222+
223+
fn sum_all(table: pgx::SpiTupleTable) -> i32 {
224+
table.map(|r| r.by_ordinal(1).unwrap().value::<i32>().unwrap()).sum()
225+
}
226+
assert_eq!(sum_all(portal.fetch(3)), 1 + 2 + 3);
227+
assert_eq!(sum_all(portal.fetch(3)), 4 + 5 + 6);
228+
assert_eq!(sum_all(portal.fetch(3)), 7 + 8 + 9);
229+
assert_eq!(sum_all(portal.fetch(3)), 10);
230+
});
231+
}
232+
212233
#[pg_test]
213234
fn test_cursor_by_name() {
214235
let cursor_name = Spi::connect(|mut client| {

pgx/src/spi.rs

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -412,6 +412,28 @@ impl<'a> SpiClient<'a> {
412412
&self,
413413
query: &str,
414414
args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
415+
) -> SpiCursor {
416+
self.open_cursor_impl(query, args)
417+
}
418+
419+
/// Set up a cursor that will execute the specified update (mutating) query
420+
///
421+
/// Rows may be then fetched using [`SpiCursor::fetch`].
422+
///
423+
/// See [`SpiCursor`] docs for usage details.
424+
pub fn open_cursor_mut(
425+
&mut self,
426+
query: &str,
427+
args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
428+
) -> SpiCursor {
429+
self.readonly = false;
430+
self.open_cursor_impl(query, args)
431+
}
432+
433+
fn open_cursor_impl(
434+
&self,
435+
query: &str,
436+
args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
415437
) -> SpiCursor {
416438
let src = std::ffi::CString::new(query).expect("query contained a null byte");
417439
let args = args.unwrap_or_default();
@@ -447,7 +469,7 @@ impl<'a> SpiClient<'a> {
447469
argtypes.as_mut_ptr(),
448470
datums.as_mut_ptr(),
449471
nulls.as_ptr(),
450-
false,
472+
self.readonly,
451473
0,
452474
)
453475
})

0 commit comments

Comments
 (0)