Fix flaky connection to redshift data API (#1834)

achals · web-flow · commit fda5b55e91c0 · 2021-09-03T13:03:10.000-07:00
* Fix flaky connection to redshift data API

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;

* Lower max

Signed-off-by: Achal Shah &lt;achals@gmail.com&gt;
diff --git a/sdk/python/feast/infra/utils/aws_utils.py b/sdk/python/feast/infra/utils/aws_utils.py
@@ -15,7 +15,7 @@
 try:
     import boto3
     from botocore.config import Config
-    from botocore.exceptions import ClientError
+    from botocore.exceptions import ClientError, ConnectionClosedError
 except ImportError as e:
     from feast.errors import FeastExtrasDependencyImportError
 
@@ -50,6 +50,10 @@ def get_bucket_and_key(s3_path: str) -> Tuple[str, str]:
     return bucket, key
 
 
+@retry(
+    wait=wait_exponential(multiplier=1, max=4),
+    retry=retry_if_exception_type(ConnectionClosedError),
+)
 def execute_redshift_statement_async(
     redshift_data_client, cluster_id: str, database: str, user: str, query: str
 ) -> dict:
