DH-4653 Add support for AWS Athena (#164)

Dataherald · May 7, 2024 · 5d6d458 · 5d6d458
1 parent 6b455be
commit 5d6d458
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -192,7 +192,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
 
@@ -257,6 +257,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
@@ -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
@@ -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