Skip to content

Commit

Permalink
Set a 5 minute limit for redshift statement execution
Browse files Browse the repository at this point in the history
Signed-off-by: Achal Shah <achals@gmail.com>
  • Loading branch information
achals committed Sep 30, 2021
1 parent ce5a130 commit 6c9ddcf
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion sdk/python/feast/infra/utils/aws_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down

0 comments on commit 6c9ddcf

Please sign in to comment.