title: Run Multiple Athena Queries in Airflow 2.0
body: <p>I am trying to create a DAG in which one of the task does <code>athena</code> query using <code>boto3</code>. It worked for one query however I am facing issues when I try to run multiple athena queries.</p>
  <p>This problem can be broken as follows:-</p>
  <ol>
  <li>If one goes through <a href="https://www.ilkkapeltola.fi/2018/04/simple-way-to-query-amazon-athena-in.html" rel="nofollow noreferrer">this</a> blog, it can be seen that <code>athena</code> uses <code>start_query_execution</code> to trigger query and <code>get_query_execution</code> for getting <code>status</code>, <code>queryExecutionId</code> and other data about the query (docs for <a href="https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/athena.html" rel="nofollow noreferrer">athena</a>)</li>
  </ol>
  <p>After following the above pattern I have following code:-</p>
  <pre><code>import json
  import time
  import asyncio
  import boto3
  import logging
  from airflow import DAG
  from airflow.operators.python import PythonOperator
  
  
  def execute_query(client, query, database, output_location):
      response = client.start_query_execution(
          QueryString=query,
          QueryExecutionContext={
              'Database': database
          },
          ResultConfiguration={
              'OutputLocation': output_location
          }
      )
  
      return response['QueryExecutionId']
  
  
  async def get_ids(client_athena, query, database, output_location):
      query_responses = []
      for i in range(5):
          query_responses.append(execute_query(client_athena, query, database, output_location))    
  
      res = await asyncio.gather(*query_responses, return_exceptions=True)
  
      return res
  
  def run_athena_query(query, database, output_location, region_name, **context):
      BOTO_SESSION = boto3.Session(
          aws_access_key_id = 'YOUR_KEY',
          aws_secret_access_key = 'YOUR_ACCESS_KEY')
      client_athena = BOTO_SESSION.client('athena', region_name=region_name)
  
      loop = asyncio.get_event_loop()
      query_execution_ids = loop.run_until_complete(get_ids(client_athena, query, database, output_location))
      loop.close()
  
      repetitions = 900
      error_messages = []
      s3_uris = []
  
      while repetitions &gt; 0 and len(query_execution_ids) &gt; 0:
          repetitions = repetitions - 1
          
          query_response_list = client_athena.batch_get_query_execution(
              QueryExecutionIds=query_execution_ids)['QueryExecutions']
        
          for query_response in query_response_list:
              if 'QueryExecution' in query_response and \
                      'Status' in query_response['QueryExecution'] and \
                      'State' in query_response['QueryExecution']['Status']:
                  state = query_response['QueryExecution']['Status']['State']
  
                  if state in ['FAILED', 'CANCELLED']:
                      error_reason = query_response['QueryExecution']['Status']['StateChangeReason']
                      error_message = 'Final state of Athena job is {}, query_execution_id is {}. Error: {}'.format(
                              state, query_execution_id, error_message
                          )
                      error_messages.append(error_message)
                      query_execution_ids.remove(query_response['QueryExecutionId'])
                  
                  elif state == 'SUCCEEDED':
                      result_location = query_response['QueryExecution']['ResultConfiguration']['OutputLocation']
                      s3_uris.append(result_location)
                      query_execution_ids.remove(query_response['QueryExecutionId'])
                   
                      
          time.sleep(2)
      
      logging.exception(error_messages)
      return s3_uris
  
  
  DEFAULT_ARGS = {
      'owner': 'ubuntu',
      'depends_on_past': True,
      'start_date': datetime(2021, 6, 8),
      'retries': 0,
      'concurrency': 2
  }
  
  with DAG('resync_job_dag', default_args=DEFAULT_ARGS, schedule_interval=None) as dag:
  
      ATHENA_QUERY = PythonOperator(
          task_id='athena_query',
          python_callable=run_athena_query,
          provide_context=True,
          op_kwargs={
              'query': 'SELECT request_timestamp FROM &quot;sampledb&quot;.&quot;elb_logs&quot; limit 10;', # query provide in athena tutorial
              'database':'sampledb',
              'output_location':'YOUR_BUCKET',
              'region_name':'YOUR_REGION'
          }
      )
  
      ATHENA_QUERY
  </code></pre>
  <p>On running above code, I am getting following error:-</p>
  <pre><code>[2021-06-16 20:34:52,981] {taskinstance.py:1455} ERROR - An asyncio.Future, a coroutine or an awaitable is required
  Traceback (most recent call last):
    File &quot;/home/ubuntu/venv/lib/python3.6/site-packages/airflow/models/taskinstance.py&quot;, line 1112, in _run_raw_task
      self._prepare_and_execute_task_with_callbacks(context, task)
    File &quot;/home/ubuntu/venv/lib/python3.6/site-packages/airflow/models/taskinstance.py&quot;, line 1285, in _prepare_and_execute_task_with_callbacks
      result = self._execute_task(context, task_copy)
    File &quot;/home/ubuntu/venv/lib/python3.6/site-packages/airflow/models/taskinstance.py&quot;, line 1315, in _execute_task
      result = task_copy.execute(context=context)
    File &quot;/home/ubuntu/venv/lib/python3.6/site-packages/airflow/operators/python.py&quot;, line 117, in execute
      return_value = self.execute_callable()
    File &quot;/home/ubuntu/venv/lib/python3.6/site-packages/airflow/operators/python.py&quot;, line 128, in execute_callable
      return self.python_callable(*self.op_args, **self.op_kwargs)
    File &quot;/home/ubuntu/iac-airflow/dags/helper/tasks.py&quot;, line 93, in run_athena_query
      query_execution_ids = loop.run_until_complete(get_ids(client_athena, query, database, output_location))
    File &quot;/usr/lib/python3.6/asyncio/base_events.py&quot;, line 484, in run_until_complete
      return future.result()
    File &quot;/home/ubuntu/iac-airflow/dags/helper/tasks.py&quot;, line 79, in get_ids
      res = await asyncio.gather(*query_responses, return_exceptions=True)
    File &quot;/usr/lib/python3.6/asyncio/tasks.py&quot;, line 602, in gather
      fut = ensure_future(arg, loop=loop)
    File &quot;/usr/lib/python3.6/asyncio/tasks.py&quot;, line 526, in ensure_future
      raise TypeError('An asyncio.Future, a coroutine or an awaitable is '
  TypeError: An asyncio.Future, a coroutine or an awaitable is required
  </code></pre>
  <p>I am unable to get where I am going wrong. Would appreciate some hint over the issue</p>