Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_query(self):
self.p.query('select 1')
self.assertRaises(PostgresConnectionException, self.p.query, 'RetryFailedError')
self.assertRaises(psycopg2.ProgrammingError, self.p.query, 'blabla')
def test_update_lock(self):
self.p.last_operation = Mock(side_effect=PostgresConnectionException(''))
self.assertTrue(self.ha.update_lock(True))
def test__query(self):
self.assertRaises(PostgresConnectionException, self.p._query, 'blabla')
self.p._state = 'restarting'
self.assertRaises(RetryFailedError, self.p._query, 'blabla')
def query(self, sql, *args, **kwargs):
if not kwargs.get('retry', True):
return self._query(sql, *args)
try:
return self.retry(self._query, sql, *args)
except RetryFailedError as e:
raise PostgresConnectionException(str(e))
def _cluster_info_state_get(self, name):
if not self._cluster_info_state:
try:
result = self._is_leader_retry(self._query, self.cluster_info_query).fetchone()
self._cluster_info_state = dict(zip(['timeline', 'wal_position', 'pg_control_timeline'], result))
except RetryFailedError as e: # SELECT failed two times
self._cluster_info_state = {'error': str(e)}
if not self.is_starting() and self.pg_isready() == STATE_REJECT:
self.set_state('starting')
if 'error' in self._cluster_info_state:
raise PostgresConnectionException(self._cluster_info_state['error'])
return self._cluster_info_state.get(name)
cursor = self._connection.cursor()
cursor.execute(sql, params)
return cursor
except psycopg2.Error as e:
if cursor and cursor.connection.closed == 0:
# When connected via unix socket, psycopg2 can't recoginze 'connection lost'
# and leaves `_cursor_holder.connection.closed == 0`, but psycopg2.OperationalError
# is still raised (what is correct). It doesn't make sense to continiue with existing
# connection and we will close it, to avoid its reuse by the `cursor` method.
if isinstance(e, psycopg2.OperationalError):
self._connection.close()
else:
raise e
if self.state == 'restarting':
raise RetryFailedError('cluster is being restarted')
raise PostgresConnectionException('connection problems')
cursor = self._cursor()
cursor.execute(sql, params)
return cursor
except psycopg2.Error as e:
if cursor and cursor.connection.closed == 0:
# When connected via unix socket, psycopg2 can't recoginze 'connection lost'
# and leaves `_cursor_holder.connection.closed == 0`, but psycopg2.OperationalError
# is still raised (what is correct). It doesn't make sense to continiue with existing
# connection and we will close it, to avoid its reuse by the `_cursor` method.
if isinstance(e, psycopg2.OperationalError):
self.close_connection()
else:
raise e
if self.state == 'restarting':
raise RetryFailedError('cluster is being restarted')
raise PostgresConnectionException('connection problems')