def wait_query(query_execution_id: str, boto3_session: Optional[boto3.Session] = None) -> Dict[str, Any]:
"""Wait for the query end.
Parameters
----------
query_execution_id : str
Athena query execution ID.
boto3_session : boto3.Session(), optional
Boto3 Session. The default boto3 session will be used if boto3_session receive None.
Returns
-------
Dict[str, Any]
Dictionary with the get_query_execution response.
Examples
--------
>>> import awswrangler as wr
>>> res = wr.athena.wait_query(query_execution_id='query-execution-id')
"""
client_athena: boto3.client = _utils.client(service_name="athena", session=boto3_session)
response: Dict[str, Any] = client_athena.get_query_execution(QueryExecutionId=query_execution_id)
state: str = response["QueryExecution"]["Status"]["State"]
while state not in _QUERY_FINAL_STATES:
time.sleep(_QUERY_WAIT_POLLING_DELAY)
response = client_athena.get_query_execution(QueryExecutionId=query_execution_id)
state = response["QueryExecution"]["Status"]["State"]
_logger.debug("state: %s", state)
_logger.debug("StateChangeReason: %s", response["QueryExecution"]["Status"].get("StateChangeReason"))
if state == "FAILED":
raise exceptions.QueryFailed(response["QueryExecution"]["Status"].get("StateChangeReason"))
if state == "CANCELLED":
raise exceptions.QueryCancelled(response["QueryExecution"]["Status"].get("StateChangeReason"))
return cast(Dict[str, Any], response["QueryExecution"])
def wait_query(
query_id: str,
boto3_session: Optional[boto3.Session] = None) -> Dict[str, Any]:
"""Wait query ends.
https://docs.aws.amazon.com/AmazonCloudWatch/latest/logs/CWL_QuerySyntax.html
Parameters
----------
query_id : str
Query ID.
boto3_session : boto3.Session(), optional
Boto3 Session. The default boto3 session will be used if boto3_session receive None.
Returns
-------
Dict[str, Any]
Query result payload.
Examples
--------
>>> import awswrangler as wr
>>> query_id = wr.cloudwatch.start_query(
... log_group_names=["loggroup"],
... query="fields @timestamp, @message | sort @timestamp desc | limit 5",
... )
... response = wr.cloudwatch.wait_query(query_id=query_id)
"""
final_states: List[str] = ["Complete", "Failed", "Cancelled"]
client_logs: boto3.client = _utils.client(service_name="logs",
session=boto3_session)
response: Dict[str, Any] = client_logs.get_query_results(queryId=query_id)
status: str = response["status"]
while status not in final_states:
time.sleep(_QUERY_WAIT_POLLING_DELAY)
response = client_logs.get_query_results(queryId=query_id)
status = response["status"]
_logger.debug("status: %s", status)
if status == "Failed":
raise exceptions.QueryFailed(f"query ID: {query_id}")
if status == "Cancelled":
raise exceptions.QueryCancelled(f"query ID: {query_id}")
return response
