Make retry config a bit configurable (Tomme#35)

README.md

@@ -49,6 +49,8 @@ A dbt profile can be configured to run against AWS Athena using the following co
 | 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`                 |
 | aws_profile_name| Profile to use from your AWS shared credentials file.                           | Optional   | `my-profile`        |
+| work_group| Identifier of Athena workgroup   | Optional   | `my-custom-workgroup`        |
+| num_retries| Number of times to retry a failing query | Optional  | `3`  | `5`
 
 **Example profiles.yml entry:**
 ```yaml

dbt/adapters/athena/connections.py

@@ -11,6 +11,8 @@
 from pyathena.cursor import Cursor
 from pyathena.error import ProgrammingError, OperationalError
 from pyathena.formatter import Formatter
+from pyathena.util import RetryConfig
+
 # noinspection PyProtectedMember
 from pyathena.formatter import _DEFAULT_FORMATTERS, _escape_hive, _escape_presto
 
@@ -29,10 +31,8 @@ class AthenaCredentials(Credentials):
     work_group: Optional[str] = None
     aws_profile_name: Optional[str] = None
     poll_interval: float = 1.0
-    _ALIASES = {
-        "catalog": "database"
-    }
-    # TODO Add pyathena.util.RetryConfig ?
+    _ALIASES = {"catalog": "database"}
+    num_retries: Optional[int] = 5
 
     @property
     def type(self) -> str:
@@ -116,7 +116,15 @@ def open(cls, connection: Connection) -> Connection:
                 cursor_class=AthenaCursor,
                 formatter=AthenaParameterFormatter(),
                 poll_interval=creds.poll_interval,
-                profile_name=creds.aws_profile_name
+                profile_name=creds.aws_profile_name,
+                retry_config=RetryConfig(
+                    attempt=creds.num_retries,
+                    exceptions=(
+                        "ThrottlingException",
+                        "TooManyRequestsException",
+                        "InternalServerException",
+                    ),
+                ),
             )
 
             connection.state = "open"