async def iterate(
self, query: ClauseElement
) -> typing.AsyncGenerator[typing.Any, None]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
cursor = await self._connection.cursor()
try:
await cursor.execute(query, args)
metadata = ResultMetaData(context, cursor.description)
async for row in cursor:
yield RowProxy(metadata, row, metadata._processors, metadata._keymap)
except psycopg2.InterfaceError as error:
raise exceptions.InterfaceError(str(query), None, error) from error
except psycopg2.DataError as error:
raise exceptions.DataError(str(query), None, error) from error
except psycopg2.OperationalError as error:
raise exceptions.OperationalError(str(query), None, error) from error
except psycopg2.IntegrityError as error:
raise exceptions.IntegrityError(str(query), None, error) from error
except psycopg2.InternalError as error:
raise exceptions.InternalError(str(query), None, error) from error
except psycopg2.ProgrammingError as error:
raise exceptions.ProgrammingError(str(query), None, error) from error
except psycopg2.NotSupportedError as error:
raise exceptions.NotSupportedError(str(query), None, error) from error
except psycopg2.DatabaseError as error:
raise exceptions.DatabaseError(str(query), None, error) from error
finally:
cursor.close()
async def fetch_all(self,
query: ClauseElement) -> typing.List[typing.Mapping]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
try:
async with self._connection.execute(query, args) as cursor:
rows = await cursor.fetchall()
metadata = ResultMetaData(context, cursor.description)
return [
RowProxy(metadata, row, metadata._processors,
metadata._keymap) for row in rows
]
except sqlite3.InterfaceError as error:
raise exceptions.InterfaceError(str(query), None, error) from error
except sqlite3.DataError as error:
raise exceptions.DataError(str(query), None, error) from error
except sqlite3.OperationalError as error:
raise exceptions.OperationalError(str(query), None,
error) from error
except sqlite3.IntegrityError as error:
raise exceptions.IntegrityError(str(query), None, error) from error
except sqlite3.InternalError as error:
raise exceptions.InternalError(str(query), None, error) from error
except sqlite3.ProgrammingError as error:
raise exceptions.ProgrammingError(str(query), None,
error) from error
except sqlite3.NotSupportedError as error:
raise exceptions.NotSupportedError(str(query), None,
error) from error
except sqlite3.DatabaseError as error:
raise exceptions.DatabaseError(str(query), None, error) from error
async def fetch_one(self, query: ClauseElement) -> typing.Optional[typing.Mapping]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
cursor = await self._connection.cursor()
try:
await cursor.execute(query, args)
row = await cursor.fetchone()
if row is None:
return None
metadata = ResultMetaData(context, cursor.description)
return RowProxy(metadata, row, metadata._processors, metadata._keymap)
except psycopg2.InterfaceError as error:
raise exceptions.InterfaceError(str(query), None, error) from error
except psycopg2.DataError as error:
raise exceptions.DataError(str(query), None, error) from error
except psycopg2.OperationalError as error:
raise exceptions.OperationalError(str(query), None, error) from error
except psycopg2.IntegrityError as error:
raise exceptions.IntegrityError(str(query), None, error) from error
except psycopg2.InternalError as error:
raise exceptions.InternalError(str(query), None, error) from error
except psycopg2.ProgrammingError as error:
raise exceptions.ProgrammingError(str(query), None, error) from error
except psycopg2.NotSupportedError as error:
raise exceptions.NotSupportedError(str(query), None, error) from error
except psycopg2.DatabaseError as error:
raise exceptions.DatabaseError(str(query), None, error) from error
finally:
cursor.close()
async def fetch_one(self, query: ClauseElement) -> RowProxy:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
async with self._connection.execute(query, args) as cursor:
row = await cursor.fetchone()
metadata = ResultMetaData(context, cursor.description)
return RowProxy(metadata, row, metadata._processors, metadata._keymap)
async def iterate(
self, query: ClauseElement
) -> typing.AsyncGenerator[typing.Any, None]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
async with self._connection.execute(query, args) as cursor:
metadata = ResultMetaData(context, cursor.description)
async for row in cursor:
yield RowProxy(metadata, row, metadata._processors, metadata._keymap)
async def fetch_all(self, query: ClauseElement) -> typing.List[typing.Mapping]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
async with self._connection.execute(query, args) as cursor:
rows = await cursor.fetchall()
metadata = ResultMetaData(context, cursor.description)
return [
RowProxy(metadata, row, metadata._processors, metadata._keymap)
for row in rows
]
def verify_presto_column(self, column, expected_results):
inspector = mock.Mock()
inspector.engine.dialect.identifier_preparer.quote_identifier = mock.Mock()
keymap = {'Column': (None, None, 0),
'Type': (None, None, 1),
'Null': (None, None, 2)}
row = RowProxy(mock.Mock(), column, [None, None, None, None], keymap)
inspector.bind.execute = mock.Mock(return_value=[row])
results = PrestoEngineSpec.get_columns(inspector, '', '')
self.assertEqual(len(expected_results), len(results))
for expected_result, result in zip(expected_results, results):
self.assertEqual(expected_result[0], result['name'])
self.assertEqual(expected_result[1], str(result['type']))
async def fetch_one(self, query: ClauseElement) -> typing.Optional[typing.Mapping]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
cursor = await self._connection.cursor()
try:
await cursor.execute(query, args)
row = await cursor.fetchone()
if row is None:
return None
metadata = ResultMetaData(context, cursor.description)
return RowProxy(metadata, row, metadata._processors, metadata._keymap)
finally:
await cursor.close()
def _rowproxy_fixture(self, keys, processors, row):
class MockMeta(object):
def __init__(self):
pass
metadata = MockMeta()
keymap = {}
for index, (keyobjs, processor, values) in \
enumerate(list(zip(keys, processors, row))):
for key in keyobjs:
keymap[key] = (processor, key, index)
keymap[index] = (processor, key, index)
return RowProxy(metadata, row, processors, keymap)
async def fetch_one(self, query: ClauseElement) -> Optional[RowProxy]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
async with await self._connection.cursor() as cursor:
if args:
await cursor.execute(query, *args)
else:
await cursor.execute(query)
row = await cursor.fetchone()
if row is None:
return None
metadata = ResultMetaData(context, cursor.description)
return RowProxy(metadata, row, metadata._processors,
metadata._keymap)
async def fetch_all(self, query: Union[ClauseElement,
str]) -> List[RowProxy]:
assert self._connection is not None, "Connection is not acquired"
query, args, context = self._compile(query)
async with self._connection.cursor() as cursor:
if args:
await cursor.execute(query, *args)
else:
await cursor.execute(query)
rows = await cursor.fetchall()
metadata: ResultMetaData = ResultMetaData(context,
cursor.description)
return [
RowProxy(metadata, row, metadata._processors, metadata._keymap)
for row in rows
]
