From 6c9ddcf026704508cac9419ec712e07e54a2cd0c Mon Sep 17 00:00:00 2001 From: Achal Shah Date: Thu, 30 Sep 2021 15:11:09 -0700 Subject: [PATCH] Set a 5 minute limit for redshift statement execution Signed-off-by: Achal Shah --- sdk/python/feast/infra/utils/aws_utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/utils/aws_utils.py b/sdk/python/feast/infra/utils/aws_utils.py index 6985a459db9..2faeb395595 100644 --- a/sdk/python/feast/infra/utils/aws_utils.py +++ b/sdk/python/feast/infra/utils/aws_utils.py @@ -7,7 +7,13 @@ import pandas as pd import pyarrow as pa import pyarrow.parquet as pq -from tenacity import retry, retry_if_exception_type, wait_exponential +from tenacity import ( + retry, + retry_if_exception_type, + stop_after_attempt, + stop_after_delay, + wait_exponential, +) from feast.errors import RedshiftCredentialsError, RedshiftQueryError from feast.type_map import pa_to_redshift_value_type @@ -53,6 +59,7 @@ def get_bucket_and_key(s3_path: str) -> Tuple[str, str]: @retry( wait=wait_exponential(multiplier=1, max=4), retry=retry_if_exception_type(ConnectionClosedError), + stop=stop_after_attempt(5), ) def execute_redshift_statement_async( redshift_data_client, cluster_id: str, database: str, user: str, query: str @@ -88,6 +95,7 @@ class RedshiftStatementNotFinishedError(Exception): @retry( wait=wait_exponential(multiplier=1, max=30), retry=retry_if_exception_type(RedshiftStatementNotFinishedError), + stop=stop_after_delay(300), # 300 seconds ) def wait_for_redshift_statement(redshift_data_client, statement: dict) -> None: """Waits for the Redshift statement to finish. Raises RedshiftQueryError if the statement didn't succeed.