Add user-friendly msgs for Athena errors. #239

Author: igorborgest

awswrangler/athena.py

@@ -514,7 +514,20 @@ def read_sql_query(  # pylint: disable=too-many-branches,too-many-locals,too-man
         boto3_session=session,
     )
     _logger.debug("query_id: %s", query_id)
-    query_response: Dict[str, Any] = wait_query(query_execution_id=query_id, boto3_session=session)
+    try:
+        query_response: Dict[str, Any] = wait_query(query_execution_id=query_id, boto3_session=session)
+    except exceptions.QueryFailed as ex:
+        if ctas_approach is True:
+            if "Column name not specified" in str(ex):
+                raise exceptions.InvalidArgumentValue(
+                    "Please, define all columns names in your query. (E.g. 'SELECT MAX(col1) AS max_col1, ...')"
+                )
+            if "Column type is unknown" in str(ex):
+                raise exceptions.InvalidArgumentValue(
+                    "Please, define all columns types in your query. "
+                    "(E.g. 'SELECT CAST(NULL AS INTEGER) AS MY_COL, ...')"
+                )
+        raise ex  # pragma: no cover
     if query_response["QueryExecution"]["Status"]["State"] in ["FAILED", "CANCELLED"]:  # pragma: no cover
         reason: str = query_response["QueryExecution"]["Status"]["StateChangeReason"]
         message_error: str = f"Query error: {reason}"

testing/test_awswrangler/test_data_lake.py

@@ -1675,3 +1675,10 @@ def test_store_metadata_partitions_sample_dataset(database, table, path, partiti
     assert df.c0.sum() * num_files == df2.c0.sum()
     assert df.c1.sum() * num_files == df2.c1.sum()
     assert df.c2.sum() * num_files == df2.c2.sum()
+
+
+def test_athena_undefined_column(database):
+    with pytest.raises(wr.exceptions.InvalidArgumentValue):
+        wr.athena.read_sql_query("SELECT 1", database)
+    with pytest.raises(wr.exceptions.InvalidArgumentValue):
+        wr.athena.read_sql_query("SELECT NULL", database)