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_arraysize_default(self, cursor):
self.assertEqual(cursor.arraysize, AthenaResultSet.DEFAULT_FETCH_SIZE)
def test_arraysize_default(self, cursor):
self.assertEqual(cursor.arraysize, AthenaResultSet.DEFAULT_FETCH_SIZE)
def close(self):
self._connection = None
self._query_execution = None
self._meta_data = None
self._rows = None
self._next_token = None
self._rownumber = None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
class AthenaPandasResultSet(AthenaResultSet):
_pattern_output_location = re.compile(r'^s3://(?P[a-zA-Z0-9.\-_]+)/(?P.+)$')
_converters = {
'decimal': Decimal,
'varbinary': lambda b: binascii.a2b_hex(''.join(b.split(' '))),
'json': json.loads,
}
_parse_dates = [
'date',
'time',
'time with time zone',
'timestamp',
'timestamp with time zone',
]
def __init__(self, connection, converter, query_execution, arraysize, retry_config):
def __init__(self, connection, converter, query_execution, arraysize, retry_config):
super(AthenaResultSet, self).__init__(arraysize=arraysize)
self._connection = connection
self._converter = converter
self._query_execution = query_execution
assert self._query_execution, 'Required argument `query_execution` not found.'
self._retry_config = retry_config
self._meta_data = None
self._rows = collections.deque()
self._next_token = None
if self._query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
self._rownumber = 0
self._pre_fetch()
def _collect_result_set(self, query_id):
query_execution = self._poll(query_id)
return AthenaResultSet(
connection=self._connection,
converter=self._converter,
query_execution=query_execution,
arraysize=self._arraysize,
retry_config=self._retry_config)
def execute(self, operation, parameters=None, work_group=None, s3_staging_dir=None,
cache_size=0):
self._reset_state()
self._query_id = self._execute(operation,
parameters=parameters,
work_group=work_group,
s3_staging_dir=s3_staging_dir,
cache_size=cache_size)
query_execution = self._poll(self._query_id)
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
self._result_set = AthenaResultSet(
self._connection, self._converter, query_execution, self.arraysize,
self._retry_config)
else:
raise OperationalError(query_execution.state_change_reason)
return self