From 9d274602ebf4e61cae0c3566cc8a31eccf810965 Mon Sep 17 00:00:00 2001
From: Juan Carlos Jose Camacho <jc@dataherald.com>
Date: Thu, 14 Sep 2023 14:03:48 -0600
Subject: [PATCH] DH-4653 Add support for AWS Athena

---
 README.md                               |  7 ++++++-
 docs/api.create_database_connection.rst | 12 +++++++++++-
 requirements.txt                        |  1 +
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index abbfc405..0e68c63f 100644
--- a/README.md
+++ b/README.md
@@ -151,7 +151,7 @@ Once the engine is running, you will want to use it by:
 3. Querying the data in natural language
 
 ### Connecting to your data warehouses
-We currently support connections to Postgres, BigQuery, Databricks and Snowflake. You can create connections to these warehouses through the API or at application start-up using the envars.
+We currently support connections to Postgres, BigQuery, Databricks, Snowflake and AWS Athena. You can create connections to these warehouses through the API or at application start-up using the envars.
 
 #### Connecting through the API
 
@@ -216,6 +216,11 @@ You can generate the `connection_uri` parameter in the API call for each of the
 "connection_uri": snowflake://<user>:<password>@<organization>-<account-name>/<database>/<schema>
 ```
 
+**AWS Athena**
+```
+"connection_uri": awsathena+rest://<aws_access_key_id>:<aws_secret_access_key>@athena.<region_name>.amazonaws.com:443/<schema_name>?s3_staging_dir=<s3_staging_dir>&work_group=primary
+```
+
 **BigQuery**
 To connect to BigQuery you should create a json credential file. Please follow Steps 1-3 under "Configure 
 BigQuery Authentication in Google Cloud Platform" in 
diff --git a/docs/api.create_database_connection.rst b/docs/api.create_database_connection.rst
index ad5c2799..3292b6dc 100644
--- a/docs/api.create_database_connection.rst
+++ b/docs/api.create_database_connection.rst
@@ -1,7 +1,7 @@
 Create a Database connection
 =============================
 
-We currently support connections to the following data warehouses: ``databricks``, ``postgresql``, ``snowflake``, ``bigquery``. All sensitive connection data 
+We currently support connections to the following data warehouses: ``databricks``, ``postgresql``, ``snowflake``, ``bigquery`` and ``AWS Athena``. All sensitive connection data
 is encrypted using the key you provide in your .env file before being stored to the application storage. 
 
 You can also specify the engine to connect to the Database through an SSH tunnel, as demonstrated in the second example below.
@@ -147,6 +147,16 @@ Example::
 
 "connection_uri": snowflake://jon:123456@foo-bar/my-database/public
 
+AWS Athena
+^^^^^^^^^^^^
+
+Uri structure::
+
+"connection_uri": awsathena+rest://<aws_access_key_id>:<aws_secret_access_key>@athena.<region_name>.amazonaws.com:443/<schema_name>?s3_staging_dir=<s3_staging_dir>&work_group=primary
+
+Example::
+
+"connection_uri": awsathena+rest://foobar:foobar@athena.us-east-2.amazonaws.com:443/db_test?s3_staging_dir=s3://my-bucket/output/&work_group=primary
 
 BigQuery
 ^^^^^^^^^^^^
diff --git a/requirements.txt b/requirements.txt
index 9e9d60f5..8adf5f61 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -33,3 +33,4 @@ sphinx==6.2.1
 sphinx-book-theme==1.0.1
 boto3==1.28.38
 botocore==1.31.38
+PyAthena==3.0.6