Support initsql statements for new connections

[feikesteenbergen]

I'd like to run a number of statements only once after establishing a connection to the database.

(I'm using PostgreSQL, but I guess this may be generically useful).

Issue #180 for example, would be relatively simple once this is .

Things I may wish to execute are:

SET timezone('UTC');
SET statement_timeout TO 120;
SET application_name TO 'logwriter';

Not sure what the interface should be, but something like allowing a Vec of statements to be executed, when building a connection with the builder?

    let pool = PgPool::builder()
        .max_size(1) // maximum number of connections in the pool
        .connect_timeout(Duration::new(5, 0))
        .init_sql(vec!["SET statement_timeout TO 120"]),
        .build("postgresql://jdoe:verysecret@localhost:51200/jdoe")
        .await?;

A comparable feature is https://tomcat.apache.org/tomcat-7.0-doc/jdbc-pool.html (search for initsql).

[mehcode]

• r2d2 supports this through connection_customizer
• Vert.x (Java) does not support this
• PGX (Go) supports this through a AfterConnect callback function on the Pool
• SQLAlchemy (Python) has a generic "event" bus for the pool which supports connect events
• node-postgres (Node.js) has the same event bus style support that Python does
• JDBC (Java) obviously supports .setInitSQL() though it does from googling only allow one statement?

---

let pool = PgPool::builder()
  .max_size(1)
  // .on_acquire() - called when a connection is checked out from the pool
  // .on_release() - called when a connection is returned to the pool
  // .on_close() - called when a connection is dropped
  .on_connect(|conn| async move {
    conn.execute("
SET timezone('UTC');
SET statement_timeout TO 120;
SET application_name TO 'logwriter';
    ").await
  })
  .build(" ... ").await?;

The most robust way here is probably just a series of on_x callbacks you can setup on PoolOptions. Thoughts?

cc @danielakhterov @abonander

[mingjunyang]

I fix this problem on my computer temporary, edit my local source code.
The application_name is very importantly for my pg monitor and metrics.

It's work well.

 let application_name = url
        .param(&"application_name")
        .unwrap_or(std::borrow::Cow::Borrowed(""));

// https://www.postgresql.org/docs/12/protocol-flow.html#id-1.10.5.7.3
async fn startup(stream: &mut PgStream, url: &Url) -> crate::Result<BackendKeyData> {
    // Defaults to postgres@.../postgres
    let username = url.username().unwrap_or(Cow::Borrowed("postgres"));
    let database = url.database().unwrap_or("postgres");
    let application_name = url
        .param(&"application_name")
        .unwrap_or(std::borrow::Cow::Borrowed(""));
    // See this doc for more runtime parameters
    // https://www.postgresql.org/docs/12/runtime-config-client.html
    let params = &[
        ("user", username.as_ref()),
        ("database", database),
        // Sets the display format for date and time values,
        // as well as the rules for interpreting ambiguous date input values.
        ("DateStyle", "ISO, MDY"),
        // Sets the time zone for displaying and interpreting time stamps.
        ("TimeZone", "UTC"),
        // Adjust postgres to return percise values for floats
        // NOTE: This is default in postgres 12+
        ("extra_float_digits", "3"),
        // Sets the client-side encoding (character set).
        ("client_encoding", "UTF-8"),
        ("application_name", &application_name),
    ];

    stream.write(StartupMessage { params });
    stream.flush().await?;

[janaakhterov]

@mehcode Do we want on_x calls to be callbacks or &'static str? My guess is most users would only want to execute some statements on the connection, so writing

|conn| async move {
    conn.execute("blah").await?;
}

is going to get pretty redundant quickly.

This does also beg the question if we should even support users doing anything else but executing statements in the on_x calls. With a call back arbitrary code can be executed which seems wrong to me.

[bjeanes]

With a call back arbitrary code can be executed which seems wrong to me.

That's a pretty compelling point to me. However, that might also be necessary in some cases. For instance, maybe you need to use some secret or the query is based on some feature-toggle that can be changed at runtime? I'm not actually advocating for a hook here, I think these edge cases are pretty exceptional, but they do come to mind!

[mehcode]

Callback is more flexible in the general case. We can always use a trait that would accept just a string to make that case simple as an enhancement.

[abonander]

The annoying part with using a trait is that the closure would need to give the argument types for the trait resolution to work, e.g.

|conn: &mut PgConnection| async move {

}

We encountered this already with the MapRow trait (which I think is gone in 0.4 fortunately).

[janaakhterov]

@bjeanes

maybe you need to use some secret or the query is based on some feature-toggle that can be changed at runtime?
Wouldn't this require an entirely new pool to be created since all the connections within the pool would need to update/run the query to enable/disable feature? Unless I'm misunderstanding.

@mehcode Callbacks are definitely more flexible, but I'm just thinking what can a user possibly want to do other than call some initialization queries. This arbitrary code would run on every single connection startup, not just a pool startup which is why it seems so weird.

[bjeanes]

Wouldn't this require an entirely new pool to be created since all the connections within the pool would need to update/run the query to enable/disable feature? Unless I'm misunderstanding.

This is too context-specific. I am hypothesising reasons one might need a closure, not citing a need I have. If I were in a situation like that, I'd probably make sure I killed the pool and re-established it or relied on max_lifetime option of the pool or something. The larger point I was making is: we might need to consider that it is possible to have a need to run a query when a connection is established which is not static (by which I don't just mean 'static), in which case a closure would be useful.

[bjeanes]

The annoying part with using a trait is that the closure would need to give the argument types for the trait resolution to work, e.g.

Having two separate methods doesn't seem overly onerous to me, and there's plenty of precedent in std like this.

[abonander]

So we thinking like:

impl<DB: Database> Builder<DB> {
    pub fn init_sql(&mut self, sql: Into<String>) -> &mut Self { ... }

    pub fn on_connect<F, Fut>(&mut self, on_connect: F) -> &mut Self
    where 
        F: Send + Sync + 'static + Fn(&mut DB::Connection) -> Fut, 
        Fut: Future<Output = sqlx::Result<()>> + Send 
    { ... }
}

[abonander]

With #430 we might also want an on_connect_sync() that just takes a straight synchronous callback so returning a future isn't necessary, although it could be trivially implemented in terms of on_connect().

[mehcode]

I think a first, minimal implementation that would help a lot of people is just:

impl<DB: Database> Builder<DB> {
    pub fn on_connect<F, Fut>(&mut self, f: F) -> &mut Self
    where 
        F: Send + Sync + 'static + Fn(&mut DB::Connection) -> Fut, 
        Fut: Future<Output = Result<(), Box<dyn Error + Send + Sync + 'static>>> + Send 
    { ... }

    pub fn on_close<F, Fut>(&mut self, f: F) -> &mut Self
    where 
        F: Send + Sync + 'static + Fn(&mut DB::Connection) -> Fut, 
        Fut: Future<Output = Result<(), Box<dyn Error + Send + Sync + 'static>>> + Send 
    { ... }
}

Most else is written fairly trivially in this.

For example, given #430:

let pool = Pool::builder()
    .on_connect(|conn| async move {
        unsafe {
            sqlite3_do_something(conn.as_raw_handle());
        }

        Ok(())
    })
    .build().await?;

init_sql is even simpler as:

let pool = Pool::builder()
    .on_connect(|conn| conn.execute(r#"
SQL
GOES
HERE
    "#))
    .build().await?;

[abonander]

If the error type is Box<Error + [...]> then how do we propagate that out to the caller? A new variant like sqlx::Error::PoolOnConnect(Box<Error + ...>)?

[mehcode]

That sounds right. The issue with using sqlx::Result is the user couldn't return their own errors ( like in the case of #430 where they may have a custom error when dealing with libsqlite3 directly ).

[agentsim]

@mehcode Now that PoolOptions has after_connect and after_release can this issue be closed?

[mehcode]

Yep this is now supported on master.