diff --git a/providers/amazon/tests/system/amazon/aws/example_bedrock_retrieve_and_generate.py b/providers/amazon/tests/system/amazon/aws/example_bedrock_retrieve_and_generate.py index 7bfb10ae849b8..12cdf02e37f14 100644 --- a/providers/amazon/tests/system/amazon/aws/example_bedrock_retrieve_and_generate.py +++ b/providers/amazon/tests/system/amazon/aws/example_bedrock_retrieve_and_generate.py @@ -32,7 +32,9 @@ AWSV4SignerAuth, OpenSearch, RequestsHttpConnection, + TransportError, ) +from tenacity import before_sleep_log, retry, retry_if_exception_type, stop_after_attempt, wait_fixed from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook from airflow.providers.amazon.aws.hooks.opensearch_serverless import OpenSearchServerlessHook @@ -272,6 +274,23 @@ def create_vector_index(index_name: str, collection_id: str, region: str): else: raise + # An acknowledged index-creation call does not guarantee the index is ready for use; + # OpenSearch Serverless materializes it asynchronously and exposes no status API for it. + # Poll a trivial search as a readiness probe so downstream Bedrock ingestion does not + # start (and fail) before the index can serve requests. + @retry( + retry=retry_if_exception_type(TransportError), + stop=stop_after_attempt(30), + wait=wait_fixed(5), + before_sleep=before_sleep_log(log, logging.INFO), + reraise=True, + ) + def _wait_for_index_readiness(): + oss_client.search(index=index_name, body={"query": {"match_all": {}}, "size": 0}) + + _wait_for_index_readiness() + log.info("Index %s is ready.", index_name) + @task def copy_data_to_s3(bucket: str):