How do I handle errors and retry in PyAthena?

Question

I have an Athena query that I run every day from my local Ubuntu machine. It runs fine most times.

def get_athena_data(**kwargs):
   athena_conn = connect(aws_access_key_id = access_key, aws_secret_access_key = s_key, s3_staging_dir = path, region_name = region)
   print(f"{datetime.today().strftime('%Y-%m-%d %H:%M.%S')} Athena connection established; starting to query data using pd-sql integration")
   load_data = pd.read_sql(sql,athena_conn)
   return load_data

However, the other day I got (this was miles long, so I used SNIP a few times):

Traceback (most recent call last):
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pandas/io/sql.py", line 1586, in execute
    cur.execute(*args, **kwargs)
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pyathena/util.py", line 306, in _wrapper
    return wrapped(*args, **kwargs)
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pyathena/cursor.py", line 79, in execute
    raise OperationalError(query_execution.state_change_reason)
pyathena.error.OperationalError: GENERIC_INTERNAL_ERROR: Unable to create class ... SNIP ...
]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pandas/io/sql.py", line 1590, in execute
    self.con.rollback()
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pyathena/connection.py", line 184, in rollback
    raise NotSupportedError
pyathena.error.NotSupportedError

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/ken/Documents/projects/site_alerts/code/Site_alerts_v18.py", line 174, in <module>
    sql_results = get_athena_data(SNIP)
  File "/home/ken/Documents/projects/site_alerts/code/site_alert_functions_v18.py", line 351, in get_athena_data
    load_data = pd.read_sql(sql,athena_conn)
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pandas/io/sql.py", line 412, in read_sql
    chunksize=chunksize,
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pandas/io/sql.py", line 1633, in read_query
    cursor = self.execute(*args)
  File "/home/ken/anaconda3/lib/python3.7/site-packages/pandas/io/sql.py", line 1595, in execute
    raise ex from inner_exc
pandas.io.sql.DatabaseError: Execution failed on sql: ... SNIP ... 


GENERIC_INTERNAL_ERROR: Unable to create class ... SNIP ...
]
unable to rollback

Fine, so I need to handle errors, and I'd like to retry if it does fail. So I tried:

def retry(func, max_tries=5):
    for i in range(max_tries):
        try:
           func()
           print('completed successfully')
           break
        except Exception:
           print('error')
           continue

retry(get_athena_data(ARGS))

But, that isn't working. It still halts execution when Athena fails (I fed in a flawed sql query to simulate).

How can I handle exceptions and execute retry?

I found this in the Pyathena issues but it makes no sense to me and has no usage instructions.

Answer 1

You are calling the function get_athena_data and passing its return to the function retry, not the function.

Try it this way: retry(get_athena_data).

(UPDATED) Now passing some args:

def retry(func, max_tries=5, *args, **kwargs):
    for i in range(max_tries):
        try:
           func(*args, **kwargs)
           print('completed successfully')
           break
        except Exception:
           print('error')
           continue

retry(get_athena_data, arg1, arg2, kwarg1="foo", kwarg2="bar")