Issue: Set statement_timeout per connection or per transaction
Author: tuukkamustonen
State: closed
Created: 2019-04-08 11:23:06

Using peewee with postgres, you can normally set `statement_timeout` like this:

```python
conn = PostgresqlExtDatabase(
    ...
    options='-c statement_timeout=1ms'
)
```

However, `pgbouncer` doesn't allow `options` (you can ignore it, but it just doesn't have effect then). So I'm planning to do something like:

```sql
database.connect()
database.execute_sql("SET statement_timeout = '1ms'")
```

This seems to work, but I wonder if it's the correct approach(?)

Also, when you run `pgbouncer` in `transaction` mode, each transaction may get served a different connection. In this case, I think I should rather:

```sql
BEGIN
SET LOCAL statement_timeout = '1ms'
...
```

Is it possible to achieve this with peewee so that I wouldn't need to manually do something like this:

```python
with db.atomic():
    db.execute_sql('SET ...')
```

(`transaction` pooling mode is better than `session` mode because connections don't lie idle when connection is open but no DB activity is happening.)

Peewee uses a special hook, `Database._initialize_connection(self, conn)`, which you can subclass and override to implement special logic when a new connection is opened.

Untested, but presumably something like this:

```python
class PgBouncerPostgresqlDatabase(PostgresqlDatabase):
    def _initialize_connection(self, conn):
        curs = conn.cursor()
        curs.execute('set statement_timeout=\'1ms\'')
```

I'm not sure what the ramifications of transaction-mode would be, and how that would interact with Peewee's own connection management or whatever. But yeah, I guess you could override the `begin()` method as well:

```python
class PgBouncerPostgresqlDatabase(PostgresqlDatabase):
    def _initialize_connection(self, conn):
        curs = conn.cursor()
        curs.execute('set statement_timeout=\'1ms\'')

    def begin(self):
        self.execute_sql('SET ...')
```

I'll give those a try, thanks!

Tried the suggestions above.

```python
    def begin(self):
        self.execute_sql('SET ...')
```

This seems to work only when transaction is explicitly defined (e.g. `with db.atomic()`). It doesn't trigger with `autocommit=True`.

Also, it only worked for me if I _didn't_ use `LOCAL` and only did `SET statement_timeout=...`.

Postgres docs say:

> The effects of SET LOCAL last only till the end of the current transaction, whether committed or not.--

In the logs, I cannot see any `BEGIN`, `COMMIT` or `SAVEPOINT` messages so I'm not sure what `db.atomic()` actually triggers. I guess `SET LOCAL` gets run in wrong context and thus doesn't go into effect. I looked into code but it would take time to get into it...

On the other hand:

```python
    def _initialize_connection(self, conn):
        curs = conn.cursor()
        curs.execute('set statement_timeout=\'1ms\'')
```

This works fine. So, I think I'm just going to use session-wide parameters instead of in-transaction (`LOCAL`).

> This seems to work only when transaction is explicitly defined (e.g. with db.atomic()). It doesn't trigger with autocommit=True.

Yes, that's correct.

> In the logs, I cannot see any BEGIN, COMMIT or SAVEPOINT messages so I'm not sure what db.atomic() actually triggers

The transactions are managed behind-the-scenes by `psycopg2`. You can see in the peewee code that it calls `rollback()` or `commit()` on the underlying psycopg2 connection object. The [managing transactions](http://docs.peewee-orm.com/en/latest/peewee/database.html#managing-transactions) document is worth reading if you're unclear on Peewee's APIs or how to use them.

Peewee always runs in autocommit mode by default. That is: Peewee will issue commit unless there is an active manual_commit or atomic block wrapping the code. If you want to explicitly manage transactions/commit yourself, you wrap the corresponding code in the `manual_commit()` context manager. If you want to wrap multiple operations in a transaction or savepoint, you wrap the corresponding code in an `atomic()` context manager. Otherwise **each statement is effectively in its own transaction**.

This might also help: http://initd.org/psycopg/docs/usage.html#transactions-control

A few findings, should anyone read this.

```python
    def _initialize_connection(self, conn):
        curs = conn.cursor()
        curs.execute('SET statement_timeout=\'1s\'')
```

This opens up new transaction, because:

> In Psycopg transactions are handled by the connection class. **By default, the first time a command is sent to the database (using one of the cursors created by the connection), a new transaction is created.** The following database commands will be executed in the context of the same transaction – not only the commands issued by the first cursor, but the ones issued by all the cursors created by the same connection. Should any command fail, the transaction will be aborted and no further command will be executed until a call to the rollback() method.

So, if you use the code snippet above, it opens up transaction, but does not commit. This means connection is left hanging as "idle in transaction" state. If you have slow processing or I/O, that means the connection is kept open / reserved for the whole of this time (as in my case, where I am calling another service that may sometimes take 10-30s to respond).

If you're not using connection pooler, then this is just fine. You reserve DB connections for the duration of connect-close cycle, anyway.

However, I am running pgBouncer in *transaction*-pooling mode, so each connection gets served a connection to pgBouncer, but only transactions are routed to the actual PG database, so it's important when you `BEGIN`/`COMMIT`.  Opening, but not committing the transaction shows as "idle in transaction" in PG and that transaction is then holding an actual pgBouncer->DB connection, just sitting idle, blocking others.

So commit even this statement:

```python
    def _initialize_connection(self, conn):
        with conn.cursor() as curs:  # auto-releases cursor, not required but nicer
            curs.execute('SET statement_timeout=\'1s\'')
        conn.commit()  # REQUIRED
```

Finally, with *transaction*-pooling mode in pgBouncer, nothing actually guarantees that the `SET ...` statement would be run for each *pgBouncer->PG* connection as the first call. Any transaction may get served with pgBouncer->PG connection, which means any `SELECT`, `UPDATE` or whatever you might have there may get be served a fresh pgBouncer->PG connection, where `SET ...` wasn't actually called yet.

And sure, `SET ...` statements leak from connection to connection, so it's completely messed up (but for me _that_ is fine as I use same `statement_timeout` for all connections, so leaking does no harm).

Due to above, I tried switching to *session*-pooling mode, but faced into some potential issue with pgBouncer, ref https://github.com/pgbouncer/pgbouncer/issues/384. However, even if I managed to switch to session-pooling, it wouldn't help, because I'm opening DB connection (via peewee) when inbound (web app) request processing starts and closing it when the response is about to get returned. With session-pooling mode, I would need finer-grained connect/close cycle, again to avoid connections hanging in "idle in transaction" state, doing nothing (as they wait processing or I/O).

So, potential solutions here:

1. Wrap all calls `with db.atomic()` to explicitly define transactions and use the suggested:

       class PgBouncerPostgresqlDatabase(PostgresqlDatabase):
           def begin(self):
               self.execute_sql('SET ...')

   However, wrapping calls in `with db.atomic()` blocks is problematic, because of code re-use and structure - you don't always know if a transaction is already open or not, so you might actually end up with nested transactions).

2. Write context-manager (similar to `db.atomic()`) that connects/closes connection and wrap it over all statements (this would allow to use session-pooling... but there's no real benefit here).

3. Use pgBouncer's `query_timeout` as substitute for `statement_timeout` (if it's ok, see below *)

4. Use server -level `statement_timeout` (and override it when/if needed, e.g. when debugging and running potentially slow queries)

5. Set `statement_timeout` in pgBouncer->DB connection via `connect_query` (`options=-c statement_timeout=...` would do the same but it's not allowed). I will go with this.

*pgBouncer has `query_timeout` (broken in older versions and fixed in 1.9.0+), but it's documented as:

> **query_timeout**
> Queries running longer than that are canceled. This should be used **only with slightly smaller server-side `statement_timeout`**, to apply only for network problems. [seconds]

So apparently it cannot be used as a substitute for `statement_timeout` (but unfortunately, it docs don't explain why).

Thank you for the excellent comment and sharing ur experience w/this stuff.