def pg_wait(conn, env, timeout=0):
while 1:
state = conn.poll()
if state == psycopg2.extensions.POLL_OK:
raise StopIteration
elif state == psycopg2.extensions.POLL_WRITE:
yield env['x-wsgiorg.fdevent.writable'](conn.fileno(), timeout)
elif state == psycopg2.extensions.POLL_READ:
yield env['x-wsgiorg.fdevent.readable'](conn.fileno(), timeout)
else:
raise psycopg2.OperationalError("poll() returned %s" % state)
def wait(conn):
while True:
state = conn.poll()
if state == psycopg2.extensions.POLL_OK:
break
elif state == psycopg2.extensions.POLL_WRITE:
select.select([], [conn.fileno()], [])
elif state == psycopg2.extensions.POLL_READ:
select.select([conn.fileno()], [], [])
else:
raise psycopg2.OperationalError("poll() returned %s" % state)
def mock_function(connection):
assert isinstance(connection, MockConnection)
connection.transaction_status = \
psycopg2.extensions.TRANSACTION_STATUS_INTRANS
_function_calls.append(connection)
# the default sleep times are going to be,
# 2, 4, 6, 10, 15
# so after 2 + 4 + 6 + 10 + 15 seconds
# all will be exhausted
if sum(_sleep_count) < sum([2, 4, 6, 10, 15]):
raise psycopg2.OperationalError('Arh!')
def gevent_wait_callback(connection, timeout=None):
while True:
state = connection.poll()
if state == extensions.POLL_OK:
break
elif state == extensions.POLL_READ:
gevent.socket.wait_read(connection.fileno(), timeout=timeout)
elif state == extensions.POLL_WRITE:
gevent.socket.wait_write(connection.fileno(), timeout=timeout)
else:
raise psycopg2.OperationalError("Unexcepted state: %s" % state)
def wait(conn):
while 1:
state = conn.poll()
if state == pg.extensions.POLL_OK:
break
elif state == pg.extensions.POLL_WRITE:
select.select([], [conn.fileno()], [])
elif state == pg.extensions.POLL_READ:
select.select([conn.fileno()], [], [])
else:
raise pg.OperationalError("poll() returned {}".format(state))
async def test_postgres_cursor_failure_concurrency(self, cursor):
cursor.side_effect = psycopg2.OperationalError()
def on_error(*args):
return RuntimeError
async def invoke_cursor():
async with self.app.postgres_connector(on_error) as connector:
await connector.execute('SELECT 1')
with self.assertRaises(RuntimeError):
await asyncio.gather(invoke_cursor(), invoke_cursor())
def execute(self, sql, *params):
if sql.startswith('blabla'):
raise psycopg2.ProgrammingError()
elif sql == 'CHECKPOINT' or sql.startswith(
'SELECT pg_catalog.pg_create_'):
raise psycopg2.OperationalError()
elif sql.startswith('RetryFailedError'):
raise RetryFailedError('retry')
elif sql.startswith('SELECT catalog_xmin'):
self.results = [(100, 501)]
elif sql.startswith('SELECT slot_name, catalog_xmin'):
self.results = [('ls', 100, 500, b'123456')]
elif sql.startswith('SELECT slot_name'):
self.results = [('blabla', 'physical'), ('foobar', 'physical'),
('ls', 'logical', 'a', 'b', 5, 100, 500)]
elif sql.startswith('SELECT CASE WHEN pg_catalog.pg_is_in_recovery()'):
self.results = [(1, 2, 1, 0, False, 1, 1, None, None, [{
"slot_name":
"ls",
"confirmed_flush_lsn":
12345
}])]
elif sql.startswith('SELECT pg_catalog.pg_is_in_recovery()'):
self.results = [(False, 2)]
elif sql.startswith('SELECT pg_catalog.pg_postmaster_start_time'):
replication_info = '[{"application_name":"walreceiver","client_addr":"1.2.3.4",' +\
'"state":"streaming","sync_state":"async","sync_priority":0}]'
now = datetime.datetime.now(tzutc)
self.results = [(now, 0, '', 0, '', False, now, replication_info)]
elif sql.startswith('SELECT name, setting'):
self.results = [
('wal_segment_size', '2048', '8kB', 'integer', 'internal'),
('wal_block_size', '8192', None, 'integer', 'internal'),
('shared_buffers', '16384', '8kB', 'integer', 'postmaster'),
('wal_buffers', '-1', '8kB', 'integer', 'postmaster'),
('search_path', 'public', None, 'string', 'user'),
('port', '5433', None, 'integer', 'postmaster'),
('listen_addresses', '*', None, 'string', 'postmaster'),
('autovacuum', 'on', None, 'bool', 'sighup'),
('unix_socket_directories', '/tmp', None, 'string',
'postmaster')
]
elif sql.startswith('IDENTIFY_SYSTEM'):
self.results = [('1', 3, '0/402EEC0', '')]
elif sql.startswith('TIMELINE_HISTORY '):
self.results = [('',
b'x\t0/40159C0\tno recovery target specified\n\n'
b'1\t0/40159C0\tno recovery target specified\n\n'
b'2\t0/402DD98\tno recovery target specified\n\n'
b'3\t0/403DD98\tno recovery target specified\n')]
else:
self.results = [(None, None, None, None, None, None, None, None,
None, None)]
def __wait_for_completion(cls, client):
while client is not None:
state = client.poll()
if state == psycopg2.extensions.POLL_OK:
break
elif state == psycopg2.extensions.POLL_WRITE:
select.select([], [client.fileno()], [])
elif state == psycopg2.extensions.POLL_READ:
select.select([client.fileno()], [], [])
else:
raise psycopg2.OperationalError(
"__wait_for_completion->poll() returned {}".format(state))
def test_database_connection():
try:
request_db("SELECT url, status_code, response_time, "
"regexp_pattern_found "
"FROM METRICS;")
db_version = request_db("SELECT VERSION()")
assert len(db_version) == 1
except psycopg2.OperationalError as e:
raise psycopg2.OperationalError("Database connection error") from e
except psycopg2.errors.UndefinedTable as e:
raise psycopg2.errors.UndefinedTable("Table metrics doesn't exist - "
"create table") from e
def eventlet_wait_callback(conn, timeout=-1):
"""A wait callback useful to allow eventlet to work with Psycopg."""
while 1:
state = conn.poll()
if state == extensions.POLL_OK:
break
elif state == extensions.POLL_READ:
trampoline(conn.fileno(), read=True)
elif state == extensions.POLL_WRITE:
trampoline(conn.fileno(), write=True)
else:
raise psycopg2.OperationalError("Bad result from poll: %r" % state)
def wait(conn):
"""wait for some event on connection to postgres"""
while 1:
state = conn.poll()
if state == psycopg2.extensions.POLL_OK:
break
elif state == psycopg2.extensions.POLL_WRITE:
select.select([], [conn.fileno()], [])
elif state == psycopg2.extensions.POLL_READ:
select.select([conn.fileno()], [], [])
else:
raise psycopg2.OperationalError("poll() returned %s" % state)
def gevent_wait_callback(conn, timeout=None):
"""A wait callback useful to allow gevent to work with Psycopg."""
while 1:
state = conn.poll()
if state == extensions.POLL_OK:
break
elif state == extensions.POLL_READ:
wait_read(conn.fileno(), timeout=timeout)
elif state == extensions.POLL_WRITE:
wait_write(conn.fileno(), timeout=timeout)
else:
raise psycopg2.OperationalError("Bad result from poll: %r" % state)
def wait_callback(conn, timeout=None):
while True:
state = conn.poll()
if state == extensions.POLL_OK:
return
elif state == extensions.POLL_READ:
socket.wait_read(conn.fileno(), timeout=timeout)
elif state == extensions.POLL_WRITE:
socket.wait_write(conn.fileno(), timeout=timeout)
else:
log.error('Psycopg2 driver error. Bad result')
raise psycopg2.OperationalError("Bad result from poll: %r" % state)
def psycopg2_wait_callback(conn): # pragma: no cover
fd = conn.fileno()
while True:
state = conn.poll()
if state == POLL_OK:
break
elif state == POLL_READ:
wait_to_read(fd)
elif state == POLL_WRITE:
wait_to_write(fd)
else:
raise psycopg2.OperationalError(
"Bad result from poll: %r" % state)
def _io_callback(self, future, result, fd=None, events=None):
try:
state = self.connection.poll()
except (psycopg2.Warning, psycopg2.Error) as error:
self.ioloop.remove_handler(self.fileno)
future.set_exc_info(sys.exc_info())
else:
try:
if state == POLL_OK:
self.ioloop.remove_handler(self.fileno)
future.set_result(result)
elif state == POLL_READ:
self.ioloop.update_handler(self.fileno, IOLoop.READ)
elif state == POLL_WRITE:
self.ioloop.update_handler(self.fileno, IOLoop.WRITE)
else:
future.set_exception(psycopg2.OperationalError("poll() returned %s" % state))
except IOError:
# Can happen when there are quite a lof of outstanding
# requests. See https://github.com/FSX/momoko/issues/127
self.ioloop.remove_handler(self.fileno)
future.set_exception(psycopg2.OperationalError("IOError on socker"))
def _close(self):
"""Remove the connection from the event_loop and close it."""
# N.B. If connection contains uncommitted transaction the
# transaction will be discarded
if self._fileno is not None:
self._loop.remove_reader(self._fileno)
if self._writing:
self._writing = False
self._loop.remove_writer(self._fileno)
self._conn.close()
if self._waiter is not None and not self._waiter.done():
self._waiter.set_exception(
psycopg2.OperationalError("Connection closed"))
def __init__(self):
''' Create SQL connection '''
try:
self.connection = psycopg2.connect(database=DB_NAME,
user=DB_USER,
password=DB_PASSWORD,
host=DB_HOST,
port=DB_PORT)
self.connected = True
except psycopg2.OperationalError:
logger.warning("Failed to connect to database")
self.connected = False
raise psycopg2.OperationalError("Failed to connect to database")
async def mock_cursor(self,
name=None,
cursor_factory=None,
scrollable=None,
withhold=False,
timeout=None):
global test_postgres_cursor_oer_invocation
test_postgres_cursor_oer_invocation += 1
if test_postgres_cursor_oer_invocation == 1:
raise psycopg2.OperationalError()
async with original(self, name, cursor_factory, scrollable,
withhold, timeout) as value:
yield value
def psycopg2_wait_cb(conn, timeout=-1):
"""A wait callback to trigger a yield while waiting for the database to respond
"""
while True:
state = conn.poll()
if state == extensions.POLL_OK:
break
elif state == extensions.POLL_READ:
trampoline(conn.fileno(), READ)
elif state == extensions.POLL_WRITE:
trampoline(conn.fileno(), WRITE)
else:
raise psycopg2.OperationalError(
'Bad result from poll: {}'.format(state))
def wait_cb(conn):
"""A wait callback useful to allow eventlet to work with Psycopg."""
while 1:
if panic:
raise Exception('whatever')
state = conn.poll()
if state == extensions.POLL_OK:
break
elif state == extensions.POLL_READ:
trampoline(conn.fileno(), read=True)
elif state == extensions.POLL_WRITE:
trampoline(conn.fileno(), write=True)
else:
raise psycopg2.OperationalError("Bad result from poll: %r" % state)
def test_connect_fail(self):
psycopg2_mock = Mock()
psycopg2_mock.OperationalError = psycopg2.OperationalError
psycopg2_mock.connect.side_effect = psycopg2.OperationalError()
with patch.multiple('pgsqldoc.pgsqldoc',
psycopg2=psycopg2_mock,
output=DEFAULT) as mock:
with self.assertRaises(psycopg2.OperationalError):
Preprocessor._connect(self.preprocessor, self.options)
psycopg2_mock.assert_has_calls([call.connect(self.connect_string)])
self.assertTrue(mock['output'].called)
self.assertIsNone(self.preprocessor._con)
# logger was called 2 times: first 'Trying to connect'
# second 'Failed to connect'
self.assertEqual(self.preprocessor.logger.debug.call_count, 2)
def _on_io_events(self, fd=None, events=None):
"""Invoked by Tornado's IOLoop when there are events for the fd
:param int fd: The file descriptor for the event
:param int events: The events raised
"""
if fd not in self._connections:
LOGGER.warning('Received IO event for non-existing connection')
return
try:
self._poll_connection(fd)
except OSError as error:
self._futures[fd].set_exception(
psycopg2.OperationalError('Connection error (%s)' % error))
def ready(conn, waiter):
loop = get_running_loop()
fileno = conn.fileno()
state = conn.poll()
if state == psycopg2.extensions.POLL_OK:
loop.remove_writer(fileno)
loop.remove_reader(fileno)
waiter.set_result(None)
elif state == psycopg2.extensions.POLL_WRITE:
loop.add_writer(fileno, ready, conn, waiter)
elif state == psycopg2.extensions.POLL_READ:
loop.remove_writer(fileno)
loop.add_reader(fileno, ready, conn, waiter)
else:
raise psycopg2.OperationalError("poll() returned %s" % state)
def test_basic_postgres_usage(self):
class Sneak(ConnectionContext):
def connection(self, __=None):
assert self.dsn
return MockConnection(self.dsn)
definition = Namespace()
local_config = {
'database_hostname': 'host',
'database_name': 'name',
'database_port': 'port',
'database_username': '******',
'database_password': '******',
}
postgres = Sneak(definition, local_config)
with postgres() as connection:
self.assertTrue(isinstance(connection, MockConnection))
self.assertEqual(
connection.dsn,
'host=host dbname=name port=port user=user password=password')
self.assertEqual(_closes, 0)
# exiting the context would lastly call 'connection.close()'
self.assertEqual(_closes, 1)
self.assertEqual(_commits, 0)
self.assertEqual(_rollbacks, 0)
try:
with postgres() as connection:
raise NameError('crap')
except NameError:
pass
finally:
self.assertEqual(_closes, 2) # second time
self.assertEqual(_commits, 0)
self.assertEqual(_rollbacks, 0)
try:
with postgres() as connection:
connection.transaction_status = \
psycopg2.extensions.TRANSACTION_STATUS_INTRANS
raise psycopg2.OperationalError('crap!')
# OperationalError's aren't bubbled up
except psycopg2.OperationalError:
pass
self.assertEqual(_closes, 3)
self.assertEqual(_commits, 0)
self.assertEqual(_rollbacks, 0)
def _wait(connection, timeout=5):
t0 = time.time()
while 1:
state = connection.poll()
if state == psycopg2.extensions.POLL_OK:
break
elif state == psycopg2.extensions.POLL_WRITE:
select.select([], [connection.fileno()], [], 1)
elif state == psycopg2.extensions.POLL_READ:
select.select([connection.fileno()], [], [], 1)
else:
raise psycopg2.OperationalError("poll() returned %s" % state)
if timeout > 0:
if time.time() - t0 > timeout:
raise TimeoutError("Database operation timed out")
def test_failure_acquire_lock(self, psycopg2_connector):
execute_mock = mock.MagicMock()
execute_mock.execute.side_effect = psycopg2.OperationalError("Broken")
cursor_mock = mock.MagicMock()
cursor_mock.__enter__ = mock.MagicMock(return_value=execute_mock)
cursor_mock.__exit__ = mock.MagicMock(return_value=False)
conn_mock = mock.MagicMock()
conn_mock.cursor.return_value = cursor_mock
psycopg2_connector.return_value = conn_mock
c = self._create_coordinator()
c.start()
test_lock = c.get_lock(b'test-lock')
self.assertRaises(tooz.ToozError, test_lock.acquire)
def gevent_wait_callback(conn, timeout=None):
"""A wait callback useful to allow gevent to work with Psycopg."""
# Copyright (C) 2010-2012 Daniele Varrazzo <*****@*****.**>
# This function is borrowed from psycogreen module which is licensed
# under the BSD license (see in odoo/debian/copyright)
while 1:
state = conn.poll()
if state == psycopg2.extensions.POLL_OK:
break
elif state == psycopg2.extensions.POLL_READ:
wait_read(conn.fileno(), timeout=timeout)
elif state == psycopg2.extensions.POLL_WRITE:
wait_write(conn.fileno(), timeout=timeout)
else:
raise psycopg2.OperationalError("Bad result from poll: %r" %
state)
def wait_async(conn):
"""wait for the connection to be ready for next operation
this is used for a async connection, since async mode connection will
not block any function call, and after the function call, it maybe
not finish doing its work, we need to wait for the work completion
"""
while 1:
state = conn.poll()
if state == psycopg2.extensions.POLL_OK:
break
elif state == psycopg2.extensions.POLL_WRITE:
select.select([], [conn.fileno()], [])
elif state == psycopg2.extensions.POLL_READ:
select.select([conn.fileno()], [], [])
else:
raise psycopg2.OperationalError("poll() returned %s" % state)
def execute(self, query_text):
for i in xrange(1, 1000):
try:
self.cursor = self.connection.cursor()
self.cursor.execute(query_text)
self.connection.commit()
return self._result_to_string()
except psycopg2.DatabaseError as e:
self.connection.rollback()
self.cursor.close()
if "no results to fetch" in str(e):
return ""
elif "could not serialize access" in str(e):
pass
else:
raise e
raise psycopg2.OperationalError('Retried 1000 times!!!')
def test_query_exception(self, mock_connect):
"""Test _query() when an exception is thrown."""
logging.disable(0)
# Because of psycopg2's chained method design, we need to chain mocks...
# result of psycopg2.connect()
mock_con = mock_connect.return_value
# result of con.cursor()
mock_cur = mock_con.cursor.return_value
# result of cur.execute()
mock_cur.execute.side_effect = psycopg2.OperationalError(
'test exception')
test_query = 'SELECT count(*) from now()'
with self.assertLogs(level=logging.WARNING):
dbs = DatabaseStatus()
dbs.query(test_query)
