diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py index deb60c8976251..6f238c8b60129 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py @@ -476,10 +476,17 @@ def _create_kb(): **self.create_knowledge_base_kwargs, )["knowledgeBase"]["knowledgeBaseId"] except ClientError as error: + error_message = error.response["Error"]["Message"].lower() + is_known_retryable_message = ( + "no such index" in error_message + # It may also be that permissions haven't even propagated yet to check for the index + or "server returned 401" in error_message + or "User does not have permissions" in error_message + ) if all( [ error.response["Error"]["Code"] == "ValidationException", - "no such index" in error.response["Error"]["Message"], + is_known_retryable_message, self.wait_for_indexing, self.indexing_error_max_attempts > 0, ] 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 a96266545687a..4b4d203964134 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 @@ -27,6 +27,7 @@ import boto3 from botocore.exceptions import ClientError from opensearchpy import ( + AuthenticationException, AuthorizationException, AWSV4SignerAuth, OpenSearch, @@ -265,7 +266,7 @@ def create_vector_index(index_name: str, collection_id: str, region: str): response = oss_client.indices.create(index=index_name, body=json.dumps(index_config)) log.info("Creating index: %s.", response) break - except AuthorizationException as e: + except (AuthorizationException, AuthenticationException) as e: # Index creation can take up to a minute and there is no (apparent?) way to check the current state. log.info( "Access denied; policy permissions have likely not yet propagated, %s tries remaining.",