diff options
Diffstat (limited to 'client/utils/misc.py')
| -rw-r--r-- | client/utils/misc.py | 46 |
1 files changed, 46 insertions, 0 deletions
diff --git a/client/utils/misc.py b/client/utils/misc.py index 1d5436e..6f73998 100644 --- a/client/utils/misc.py +++ b/client/utils/misc.py @@ -1,4 +1,6 @@ import os +import psycopg2 +import psycopg2.extras import time from subprocess import call, STDOUT @@ -21,3 +23,47 @@ def run_cmd(args, env=None, cwd=None): strout.seek(0) return (retcode, strout.read(), (time.time() - start)) + + +def connect(dbname, conn, cursor, nretries = 60, delay = 1.0): + '''Try opening a connection and a cursor. If it does not succeed (e.g. + when the database is performing recovery after a crash, retry multiple + times (as specified by nretries and delay in seconds). + ''' + + # if we already have connection and a cursor, return it + if conn and cursor: + return (conn, cursor) + + # we'll try repeatedly, with delays between the attempts + i = 0 + while i < nretries: + + i += 1 + + try: + conn = psycopg2.connect('host=localhost dbname=%s' % (dbname,)) + # TODO do we actually need autocommit? + conn.autocommit = True + cursor = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) + + return (conn, cursor) + except: + # connection failure - sleep for a while, then try again + time.sleep(delay) + + return (None, None) + + +def disconnect(conn, cursor): + '''Make sure we're disconnected (but prevent exceptions)''' + + try: + cursor.close() + except: + pass + + try: + conn.close() + except: + pass |