Add poll interval option (Tomme#11)

* Add poll interval option
stevenayers · Mar 30, 2021 · 7e10614 · 7e10614
1 parent 31f1a00
commit 7e10614
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/README.md b/README.md
@@ -21,12 +21,13 @@
 
 A dbt profile can be configured to run against AWS Athena using the following configuration:
 
-| Option         	| Description                                                                    	| Required? 	| Example            	|
-|----------------	|--------------------------------------------------------------------------------	|-----------	|--------------------	|
-| s3_staging_dir 	| S3 location to store Athena query results and metadata                         	| Required  	| `s3://bucket/dbt/` 	|
-| region_name    	| AWS region of your Athena instance                                             	| Required  	| `eu-west-1`        	|
-| schema         	| Specify the schema (Athena database) to build models into (lowercase **only**) 	| Required  	| `dbt`              	|
-| database       	| Specify the database (Data catalog) to build models into (lowercase **only**)  	| Required  	| `awsdatacatalog`   	|
+| Option          | Description                                                                     | Required?  | Example             |
+|---------------- |-------------------------------------------------------------------------------- |----------- |-------------------- |
+| s3_staging_dir  | S3 location to store Athena query results and metadata                          | Required   | `s3://bucket/dbt/`  |
+| region_name     | AWS region of your Athena instance                                              | Required   | `eu-west-1`         |
+| schema          | Specify the schema (Athena database) to build models into (lowercase **only**)  | Required   | `dbt`               |
+| database        | Specify the database (Data catalog) to build models into (lowercase **only**)   | Required   | `awsdatacatalog`    |
+| poll_interval   | Interval in seconds to use for polling the status of query results in Athena    | Optional   | `5`                 |
 
 **Example profiles.yml entry:**
 ```yaml

diff --git a/dbt/adapters/athena/connections.py b/dbt/adapters/athena/connections.py
@@ -27,6 +27,7 @@ class AthenaCredentials(Credentials):
     region_name: str
     schema: str
     work_group: Optional[str]
+    poll_interval: float = 1.0
     _ALIASES = {
         "catalog": "database"
     }
@@ -37,7 +38,7 @@ def type(self) -> str:
         return "athena"
 
     def _connection_keys(self) -> Tuple[str, ...]:
-        return "s3_staging_dir", "work_group", "region_name", "database", "schema"
+        return "s3_staging_dir", "work_group", "region_name", "database", "schema", "poll_interval"
 
 
 class AthenaCursor(Cursor):
@@ -112,7 +113,8 @@ def open(cls, connection: Connection) -> Connection:
                 schema_name=creds.schema,
                 work_group=creds.work_group,
                 cursor_class=AthenaCursor,
-                formatter=AthenaParameterFormatter()
+                formatter=AthenaParameterFormatter(),
+                poll_interval=creds.poll_interval,
             )
 
             connection.state = "open"