Handling and documenting ctas_approach for custom data sources. #392

Author: igorborgest

awswrangler/athena/_read.py

@@ -561,6 +561,7 @@ def read_sql_query(
     - Does not support columns with repeated names.
     - Does not support columns with undefined data types.
     - A temporary table will be created and then deleted immediately.
+    - Does not support custom data_source/catalog_id.
 
     **2** - ctas_approach=False:
 
@@ -571,6 +572,7 @@ def read_sql_query(
     - Faster for small result sizes (less latency).
     - Does not require create/delete table permissions on Glue
     - Supports timestamp with time zone.
+    - Support custom data_source/catalog_id.
 
     CONS:
 
@@ -685,6 +687,12 @@ def read_sql_query(
     >>> scanned_bytes = df.query_metadata["Statistics"]["DataScannedInBytes"]
 
     """
+    if ctas_approach and data_source not in (None, "AwsDataCatalog"):
+        raise exceptions.InvalidArgumentCombination("Queries with ctas_approach=True (default) does not support "
+                                                    "data_source values different than None and 'AwsDataCatalog'. "
+                                                    "Please check the related tutorial for more details "
+                                                    "(https://github.com/awslabs/aws-data-wrangler/blob/master/"
+                                                    "tutorials/006%20-%20Amazon%20Athena.ipynb)")
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
 
     cache_info: _CacheInfo = _check_for_cached_results(

requirements-dev.txt

@@ -18,5 +18,6 @@ sphinx==3.2.1
 sphinx_bootstrap_theme==0.7.1
 moto==1.3.16
 jupyterlab==2.2.8
+jupyter==1.0.0
 s3fs==0.4.2
 -e .

tests/test_athena.py

@@ -788,3 +788,9 @@ def test_athena_nan_inf(glue_database, ctas_approach, data_source):
     assert df.inf.iloc[0] == np.PINF
     assert df.inf_n.iloc[0] == np.NINF
     assert df.regular.iloc[0] == 1.2
+
+
+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")

tutorials/006 - Amazon Athena.ipynb

@@ -23,6 +23,7 @@
     "         - Does not support columns with repeated names.\n",
     "         - Does not support columns with undefined data types.\n",
     "         - A temporary table will be created and then deleted immediately.\n",
+    "         - Does not support custom data_source/catalog_id.\n",
     "\n",
     "\n",
     "- **ctas_approach=False**\n",
@@ -33,6 +34,7 @@
     "        - Faster for small result sizes (less latency).\n",
     "        - Does not require create/delete table permissions on Glue\n",
     "        - Supports timestamp with time zone.\n",
+    "        - Support custom data_source/catalog_id.\n",
     "    * `CONS`:\n",
     "        - Slower (But stills faster than other libraries that uses the regular Athena API)\n",
     "        - Does not handle nested types at all."
@@ -1187,4 +1189,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}