diff options
Diffstat (limited to 'lib')
| -rw-r--r-- | lib/extensions.py | 6 | ||||
| -rw-r--r-- | lib/extras.py | 24 |
2 files changed, 30 insertions, 0 deletions
diff --git a/lib/extensions.py b/lib/extensions.py index 20f78e6..f7431d1 100644 --- a/lib/extensions.py +++ b/lib/extensions.py @@ -60,6 +60,11 @@ from _psycopg import ISQLQuote from _psycopg import QueryCanceledError, TransactionRollbackError +try: + from _psycopg import set_wait_callback +except ImportError: + pass + """Isolation level values.""" ISOLATION_LEVEL_AUTOCOMMIT = 0 ISOLATION_LEVEL_READ_COMMITTED = 1 @@ -83,6 +88,7 @@ STATUS_IN_TRANSACTION = STATUS_BEGIN POLL_OK = 0 POLL_READ = 1 POLL_WRITE = 2 +POLL_ERROR = 3 """Backend transaction status values.""" TRANSACTION_STATUS_IDLE = 0 diff --git a/lib/extras.py b/lib/extras.py index 027c587..ac784b8 100644 --- a/lib/extras.py +++ b/lib/extras.py @@ -484,4 +484,28 @@ def register_tstz_w_secs(oids=None, conn_or_curs=None): return _ext.TSTZ_W_SECS +import select +from psycopg2.extensions import POLL_OK, POLL_READ, POLL_WRITE +from psycopg2 import OperationalError + +def wait_select(conn, curs=None): + """Wait until a connection or cursor has data available. + + The function is an example of a wait callback to be registered with + `~psycopg2.extensions.set_wait_callback()`. This function uses `!select()` + to wait for data available. + """ + poll = (curs or conn).poll + while 1: + state = poll() + if state == POLL_OK: + break + elif state == POLL_READ: + select.select([conn.fileno()], [], []) + elif state == POLL_WRITE: + select.select([], [conn.fileno()], []) + else: + raise OperationalError("bad state from poll: %s" % state) + + __all__ = filter(lambda k: not k.startswith('_'), locals().keys()) |