Skip to content

Commit

Permalink
Problem: cursors may or may not be readonly-friendly
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
yrashk committed Dec 15, 2022
1 parent 49f1037 commit 01d1d24
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 1 deletion.
21 changes: 21 additions & 0 deletions pgx-tests/src/tests/spi_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -209,6 +209,27 @@ mod tests {
});
}

#[pg_test]
fn test_cursor_mut() {
Spi::execute(|mut client| {
client.update("CREATE TABLE tests.cursor_table (id int)", None, None);

let mut portal = client.open_cursor_mut(
"INSERT INTO tests.cursor_table (id) \
SELECT i FROM generate_series(1, 10) AS t(i) RETURNING id",
None,
);

fn sum_all(table: pgx::SpiTupleTable) -> i32 {
table.map(|r| r.by_ordinal(1).unwrap().value::<i32>().unwrap()).sum()
}
assert_eq!(sum_all(portal.fetch(3)), 1 + 2 + 3);
assert_eq!(sum_all(portal.fetch(3)), 4 + 5 + 6);
assert_eq!(sum_all(portal.fetch(3)), 7 + 8 + 9);
assert_eq!(sum_all(portal.fetch(3)), 10);
});
}

#[pg_test]
fn test_cursor_by_name() {
let cursor_name = Spi::connect(|mut client| {
Expand Down
24 changes: 23 additions & 1 deletion pgx/src/spi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -412,6 +412,28 @@ impl<'a> SpiClient<'a> {
&self,
query: &str,
args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
) -> SpiCursor {
self.open_cursor_impl(query, args)
}

/// Set up a cursor that will execute the specified update (mutating) query
///
/// Rows may be then fetched using [`SpiCursor::fetch`].
///
/// See [`SpiCursor`] docs for usage details.
pub fn open_cursor_mut(
&mut self,
query: &str,
args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
) -> SpiCursor {
self.readonly = false;
self.open_cursor_impl(query, args)
}

fn open_cursor_impl(
&self,
query: &str,
args: Option<Vec<(PgOid, Option<pg_sys::Datum>)>>,
) -> SpiCursor {
let src = std::ffi::CString::new(query).expect("query contained a null byte");
let args = args.unwrap_or_default();
Expand Down Expand Up @@ -447,7 +469,7 @@ impl<'a> SpiClient<'a> {
argtypes.as_mut_ptr(),
datums.as_mut_ptr(),
nulls.as_ptr(),
false,
self.readonly,
0,
)
})
Expand Down

0 comments on commit 01d1d24

Please sign in to comment.