Set SPI connection to nonatomic (by default?) to enable txn rollbacks

[olirice]

Currently, the SPI connection uses SPI_connect() to establish a connection.

https://github.com/tcdi/pgx/blob/7163fe2437f8305b57e461697005d48de3526bbb/pgx/src/spi.rs#L213-L218

By default, calls to SPI_commit and SPI_rollback on a connection opened with SPI_connect result in an immediate error (panic). That prevents us from rolling back transactions started through pgx::Spi without a panic.

From the Postgres docs there is a function SPI_connect_ext that takes an option SPI_OPT_NONATOMIC to enable SPI_commit/SPI_rollback and is otherwise equivalent to SPI_connect.

Here's the relevant excerpt:

SPI_OPT_NONATOMIC Sets the SPI connection to be nonatomic, which means that transaction control calls (SPI_commit, SPI_rollback) are allowed. Otherwise, calling those functions will result in an immediate error.

As far as I can tell, it's a strict superset of functionality. What do you think about enabling nonatomic by default on SPI connections?

[workingjubilee]

I am mildly concerned about users depending on implementation details of Spi::connect(), which it sounds like this will incur, but I am interested in considering it. Can you show an example Rust implementation that would use this functionality?

[olirice]

Sure! The real-world use case is for pg_graphql. That project has been fully ported to pgx and the only failing test is the ability to do multiple mutations in a single transaction, rolling the whole thing back if any of them fail w/ a good error message.

Here's an example of the behavior I'm aiming to enable

---

or a minimal example could be aborting an update if too many rows are impacted:

use pgx::*;
use std::panic;

pg_module_magic!();

extension_sql!(
    r#"   
create table example (
    id serial8 not null primary key,
    title text
);

insert into example(title) values ('a'), ('b');
"#,
    name = "example",
);

#[pg_extern]
fn spi_rollback_test() -> JsonB {
    let mut errors = vec![];

    Spi::connect(|mut client| {
        let updated_record_count: usize = client
            .update(
                "update example
                set title = 'abc'
                returning 1;",
                None,
                None,
            )
            .count();

        if updated_record_count > 1 {
            errors.push("Update aborted: More than 1 record".to_string());
            client.rollback(); // <-------- ROLLBACK HERE
        }

        // Optionally do other stuff with the client here

        Ok(Some(()))
    });

    JsonB {
        0: serde_json::json!({ "errors": errors }),
    }
}

[workingjubilee]

So, this function and constant you mentioned isn't defined in PG10, which will still be supported by PGX until later this year or maybe a bit into 2023.

I am okay with currently having this cfg internally and saying Spi::connect will set options dependent on Postgres and PGX version and, when we drop support for PG10, revisiting the question of whether we should promise anything regarding this on a general basis. If we need to make it non-implicit, we will likely make it configurable then through some explicit interface. It depends mostly on if the change surfaces any performance differences (I expect that making these non-atomic may significantly impact how Postgres performs... or not).

[yrashk]

Related PR (early work in progress): #714