Fix Athena read with ctas_approach=False and chunksize=True #458

Author: igorborgest

awswrangler/athena/_read.py

@@ -4,6 +4,7 @@
 import datetime
 import logging
 import re
+import sys
 import uuid
 from typing import Any, Dict, Iterator, List, Match, NamedTuple, Optional, Union
 
@@ -699,6 +700,7 @@ def read_sql_query(
             "(https://github.com/awslabs/aws-data-wrangler/blob/master/"
             "tutorials/006%20-%20Amazon%20Athena.ipynb)"
         )
+    chunksize = sys.maxsize if ctas_approach is False and chunksize is True else chunksize
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
 
     cache_info: _CacheInfo = _check_for_cached_results(

tests/test_athena.py

@@ -804,3 +804,9 @@ def test_athena_ctas_data_source(glue_database):
     sql = "SELECT nan() AS nan, infinity() as inf, -infinity() as inf_n, 1.2 as regular"
     with pytest.raises(wr.exceptions.InvalidArgumentCombination):
         wr.athena.read_sql_query(sql, glue_database, True, data_source="foo")
+
+
+def test_chunked_ctas_false(glue_database):
+    sql = "SELECT 1 AS foo, 2 AS boo"
+    df_iter = wr.athena.read_sql_query(sql, database=glue_database, ctas_approach=False, chunksize=True)
+    assert len(list(df_iter)) == 1