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_fetch_no_data(self, cursor):
self.assertRaises(ProgrammingError, cursor.fetchone)
self.assertRaises(ProgrammingError, cursor.fetchmany)
self.assertRaises(ProgrammingError, cursor.fetchall)
def __fetch(self, next_token=None):
if not self._query_execution.query_id:
raise ProgrammingError('QueryExecutionId is none or empty.')
if self._query_execution.state != AthenaQueryExecution.STATE_SUCCEEDED:
raise ProgrammingError('QueryExecutionState is not SUCCEEDED.')
request = {
'QueryExecutionId': self._query_execution.query_id,
'MaxResults': self._arraysize,
}
if next_token:
request.update({'NextToken': next_token})
try:
response = retry_api_call(self._connection.client.get_query_results,
config=self._retry_config,
logger=_logger,
**request)
except Exception as e:
_logger.exception('Failed to fetch result set.')
raise_from(OperationalError(*e.args), e)
def __fetch(self, next_token=None):
if not self._query_execution.query_id:
raise ProgrammingError('QueryExecutionId is none or empty.')
if self._query_execution.state != AthenaQueryExecution.STATE_SUCCEEDED:
raise ProgrammingError('QueryExecutionState is not SUCCEEDED.')
request = {
'QueryExecutionId': self._query_execution.query_id,
'MaxResults': self._arraysize,
}
if next_token:
request.update({'NextToken': next_token})
try:
response = retry_api_call(self._connection.client.get_query_results,
config=self._retry_config,
logger=_logger,
**request)
except Exception as e:
_logger.exception('Failed to fetch result set.')
raise_from(OperationalError(*e.args), e)
else:
return response
def format(self, operation, parameters=None):
if not operation or not operation.strip():
raise ProgrammingError('Query is none or empty.')
operation = operation.strip()
if operation.upper().startswith('SELECT') or operation.upper().startswith('WITH'):
escaper = _escape_presto
else:
escaper = _escape_hive
kwargs = dict()
if parameters:
if isinstance(parameters, dict):
for k, v in iteritems(parameters):
func = self.get_formatter(v)
kwargs.update({k: func(self, escaper, v)})
else:
raise ProgrammingError('Unsupported parameter ' +
'(Support for dict only): {0}'.format(parameters))
def arraysize(self, value):
if value <= 0 or value > self.DEFAULT_FETCH_SIZE:
raise ProgrammingError('MaxResults is more than maximum allowed length {0}.'.format(
self.DEFAULT_FETCH_SIZE))
self._arraysize = value
def _as_pandas(self):
import pandas as pd
if not self.output_location:
raise ProgrammingError('OutputLocation is none or empty.')
bucket, key = self._parse_output_location(self.output_location)
try:
response = retry_api_call(self._client.get_object,
config=self._retry_config,
logger=_logger,
Bucket=bucket,
Key=key)
except Exception as e:
_logger.exception('Failed to download csv.')
raise_from(OperationalError(*e.args), e)
else:
length = response['ContentLength']
if length:
df = pd.read_csv(io.BytesIO(response['Body'].read()),
dtype=self.dtypes,
converters=self.converters,