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