How to set statement timeout for query execution

Question

In my web app some postgres sql queries take along time for execution. I want set statement timeout for only part of them.

One part of queries must canceled by timeout, but other must work without any restriction.

In postgres exist statement_timeout function.

How to wrap SqlAlchemy query with statement_timeout function?

Like this:

SET statement_timeout TO 1000; -- timeout for one second
<sqlalchemy generated query>;
RESET statement_timeout; -- reset

Perfect way for me set timeout for query like this:

users = session.query(User).timeout(0.5).all()

SqlAlchemy must: 1) set statement timeout 2) execute query and return result 3) reset statement timeout for current session

May be other way to set timeout for query execution?

UPDATE 1. My solution

My solution is a custom connection proxy (tested with psycopg2==2.4 and SQLAlchemy==0.6.6):

from sqlalchemy.interfaces import ConnectionProxy

class TimeOutProxy(ConnectionProxy):
    def cursor_execute(self, execute, cursor, statement, parameters, context, executemany):

        timeout = context.execution_options.get('timeout', None)

        if timeout:
            c = cursor._parent.cursor()
            c.execute('SET statement_timeout TO %d;' % int(timeout * 1000))
            c.close()

        return execute(cursor, statement, parameters, context)


engine = create_engine(URL, proxy=TimeOutProxy(), pool_size=1, max_overflow=0)

This solution without reseting statement_timeout, because each SqlAlchemy query executed in isolated transaction and statement_timeout defined inside current transaction.

Usage example (timeout pаram in seconds):

Session.query(Author).execution_options(timeout=0.001).all()

Session.bind.execute(text('select * from author;') \
      .execution_options(timeout=0.001)) \
      .fetchall()

Answer 1

You should look at the extensions provided with SQLAlchemy <= 0.6:

http://www.sqlalchemy.org/docs/06/orm/interfaces.html

There are hooks where you could stick in your code for individual operations.

SQLAlchemy 0.7+ now has an event system...there might be something similar. See

http://www.sqlalchemy.org/docs/core/events.html