Peewee: Set statement_timeout per connection or per transaction

[tuukkamustonen]

Using peewee with postgres, you can normally set statement_timeout like this:

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:

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:

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:

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.)

[coleifer]

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:

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:

class PgBouncerPostgresqlDatabase(PostgresqlDatabase):
    def _initialize_connection(self, conn):
        curs = conn.cursor()
        curs.execute('set statement_timeout=\'1ms\'')

    def begin(self):
        self.execute_sql('SET ...')

[tuukkamustonen]

I'll give those a try, thanks!

[tuukkamustonen]

Tried the suggestions above.

    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:

    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).

[coleifer]

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 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.

[coleifer]

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

[tuukkamustonen]

A few findings, should anyone read this.

    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:

    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).

[coleifer]

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