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_as_pandas(self, cursor):
df = cursor.execute('SELECT * FROM one_row').as_pandas()
self.assertEqual(df.shape[0], 1)
self.assertEqual(df.shape[1], 1)
self.assertEqual([(row['number_of_rows'],) for _, row in df.iterrows()], [(1,)])
self.assertIsNotNone(cursor.query_id)
self.assertIsNotNone(cursor.query)
self.assertEqual(cursor.state, AthenaQueryExecution.STATE_SUCCEEDED)
self.assertIsNone(cursor.state_change_reason)
self.assertIsNotNone(cursor.completion_date_time)
self.assertIsInstance(cursor.completion_date_time, datetime)
self.assertIsNotNone(cursor.submission_date_time)
self.assertIsInstance(cursor.submission_date_time, datetime)
self.assertIsNotNone(cursor.data_scanned_in_bytes)
self.assertIsNotNone(cursor.execution_time_in_millis)
self.assertIsNotNone(cursor.output_location)
def test_fetchone(self, cursor):
cursor.execute('SELECT * FROM one_row')
self.assertEqual(cursor.rownumber, 0)
self.assertEqual(cursor.fetchone(), (1,))
self.assertEqual(cursor.rownumber, 1)
self.assertIsNone(cursor.fetchone())
self.assertEqual(cursor.database, SCHEMA)
self.assertIsNotNone(cursor.query_id)
self.assertIsNotNone(cursor.query)
self.assertEqual(cursor.statement_type, AthenaQueryExecution.STATEMENT_TYPE_DML)
self.assertEqual(cursor.state, AthenaQueryExecution.STATE_SUCCEEDED)
self.assertIsNone(cursor.state_change_reason)
self.assertIsNotNone(cursor.completion_date_time)
self.assertIsInstance(cursor.completion_date_time, datetime)
self.assertIsNotNone(cursor.submission_date_time)
self.assertIsInstance(cursor.submission_date_time, datetime)
self.assertIsNotNone(cursor.data_scanned_in_bytes)
self.assertIsNotNone(cursor.execution_time_in_millis)
self.assertIsNotNone(cursor.output_location)
self.assertIsNone(cursor.encryption_option)
self.assertIsNone(cursor.kms_key)
self.assertEqual(cursor.work_group, 'primary')
def test_fetchone(self, cursor):
query_id, future = cursor.execute('SELECT * FROM one_row')
result_set = future.result()
self.assertEqual(result_set.rownumber, 0)
self.assertEqual(result_set.fetchone(), (1,))
self.assertEqual(result_set.rownumber, 1)
self.assertIsNone(result_set.fetchone())
self.assertIsNotNone(result_set.query_id)
self.assertIsNotNone(result_set.query)
self.assertEqual(result_set.state, AthenaQueryExecution.STATE_SUCCEEDED)
self.assertIsNone(result_set.state_change_reason)
self.assertIsNotNone(result_set.completion_date_time)
self.assertIsInstance(result_set.completion_date_time, datetime)
self.assertIsNotNone(result_set.submission_date_time)
self.assertIsInstance(result_set.submission_date_time, datetime)
self.assertIsNotNone(result_set.data_scanned_in_bytes)
self.assertIsNotNone(result_set.execution_time_in_millis)
self.assertIsNotNone(result_set.output_location)
def test_fetchone(self, cursor):
query_id, future = cursor.execute('SELECT * FROM one_row')
result_set = future.result()
self.assertEqual(result_set.rownumber, 0)
self.assertEqual(result_set.fetchone(), (1,))
self.assertEqual(result_set.rownumber, 1)
self.assertIsNone(result_set.fetchone())
self.assertIsNotNone(result_set.query_id)
self.assertIsNotNone(result_set.query)
self.assertEqual(result_set.state, AthenaQueryExecution.STATE_SUCCEEDED)
self.assertIsNone(result_set.state_change_reason)
self.assertIsNotNone(result_set.completion_date_time)
self.assertIsInstance(result_set.completion_date_time, datetime)
self.assertIsNotNone(result_set.submission_date_time)
self.assertIsInstance(result_set.submission_date_time, datetime)
self.assertIsNotNone(result_set.data_scanned_in_bytes)
self.assertIsNotNone(result_set.execution_time_in_millis)
self.assertIsNotNone(result_set.output_location)
def test_poll(self, cursor):
query_id, _ = cursor.execute("SELECT * FROM one_row")
future = cursor.poll(query_id)
query_execution = future.result()
self.assertIn(query_execution.state, [AthenaQueryExecution.STATE_QUEUED,
AthenaQueryExecution.STATE_RUNNING,
AthenaQueryExecution.STATE_SUCCEEDED,
AthenaQueryExecution.STATE_FAILED,
AthenaQueryExecution.STATE_CANCELLED])
def test_poll(self, cursor):
query_id, _ = cursor.execute("SELECT * FROM one_row")
future = cursor.poll(query_id)
query_execution = future.result()
self.assertIn(query_execution.state, [AthenaQueryExecution.STATE_QUEUED,
AthenaQueryExecution.STATE_RUNNING,
AthenaQueryExecution.STATE_SUCCEEDED,
AthenaQueryExecution.STATE_FAILED,
AthenaQueryExecution.STATE_CANCELLED])
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
logger=_logger,
**request)
query_ids = response.get('QueryExecutionIds', None)
if not query_ids:
break # no queries left to check
next_token = response.get('NextToken', None)
query_executions = retry_api_call(
self.connection._client.batch_get_query_execution,
config=self._retry_config,
logger=_logger,
QueryExecutionIds=query_ids
).get('QueryExecutions', [])
for execution in query_executions:
if (
execution['Query'] == query and
execution['Status']['State'] == AthenaQueryExecution.STATE_SUCCEEDED and
execution['StatementType'] == AthenaQueryExecution.STATEMENT_TYPE_DML
):
query_id = execution['QueryExecutionId']
break
if query_id or next_token is None:
break
except Exception:
_logger.warning('Failed to check the cache. Moving on without cache.')
return query_id
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 = AthenaPandasResultSet(
self._connection, self._converter, query_execution, self.arraysize,
self._retry_config)
else:
raise OperationalError(query_execution.state_change_reason)
return self
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: