This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new b09ec5c7e9a Probe AOSS index readiness as the Knowledge Base role in
Bedrock RAG system test (#71745)
b09ec5c7e9a is described below
commit b09ec5c7e9a2d7b64e83ea6446196330e6d7cddc
Author: Niko Oliveira <[email protected]>
AuthorDate: Tue Aug 18 06:05:56 2026 -0700
Probe AOSS index readiness as the Knowledge Base role in Bedrock RAG system
test (#71745)
The example_bedrock_retrieve_and_generate system test intermittently
fails in the ingestion job sensor with:
[security_exception] authentication/authorization failure. Call to
Amazon OpenSearch Serverless Vector Database did not succeed.
OpenSearch Serverless data access policies propagate per principal. The
existing index readiness probe authenticates as the task's own
principal, which proves nothing about the Knowledge Base execution role
that actually performs the ingestion. When the collection activates
quickly, Bedrock can call AOSS before the policy has propagated for the
Knowledge Base role and the ingestion job fails.
Assume the Knowledge Base role and run the readiness probe with its
credentials, so the probe only passes once the ingesting principal is
authorized. The existing tenacity retry already covers the interim 403s
because AuthorizationException and AuthenticationException subclass
TransportError.
---
.../aws/example_bedrock_retrieve_and_generate.py | 36 ++++++++++++++++++++--
1 file changed, 33 insertions(+), 3 deletions(-)
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 12cdf02e37f..4a0151d5fab 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
@@ -218,13 +218,14 @@ def create_collection(collection_name: str):
@task
-def create_vector_index(index_name: str, collection_id: str, region: str):
+def create_vector_index(index_name: str, collection_id: str, region: str,
bedrock_role_arn: str):
"""
Use the OpenSearchPy client to create the vector index for the Amazon Open
Search Serverless Collection.
:param index_name: The vector index name to create.
:param collection_id: ID of the collection to be indexed.
:param region: Name of the AWS region the collection resides in.
+ :param bedrock_role_arn: Arn of the Bedrock Knowledge Base Execution Role.
"""
# Build the OpenSearch client
oss_client = OpenSearch(
@@ -278,6 +279,30 @@ def create_vector_index(index_name: str, collection_id:
str, region: str):
# 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.
+ #
+ # The probe authenticates as the Knowledge Base execution role rather than
as this task's
+ # own principal. AOSS data access policies propagate per principal, and
the Bedrock
+ # ingestion job runs as the Knowledge Base role, so probing with this
task's (already
+ # propagated) credentials would not prove the ingestion job can reach the
index. Waiting
+ # until the search succeeds as the Knowledge Base role closes the
propagation race that
+ # intermittently failed the ingestion job sensor with:
+ # [security_exception] authentication/authorization failure. Call to
Amazon OpenSearch
+ # Serverless Vector Database did not succeed.
+ assumed = StsHook().conn.assume_role(RoleArn=bedrock_role_arn,
RoleSessionName="AossIndexReadinessProbe")
+ kb_role_session = boto3.Session(
+ aws_access_key_id=assumed["Credentials"]["AccessKeyId"],
+ aws_secret_access_key=assumed["Credentials"]["SecretAccessKey"],
+ aws_session_token=assumed["Credentials"]["SessionToken"],
+ )
+ kb_role_oss_client = OpenSearch(
+ hosts=[{"host": f"{collection_id}.{region}.aoss.amazonaws.com",
"port": 443}],
+ http_auth=AWSV4SignerAuth(kb_role_session.get_credentials(), region,
"aoss"),
+ use_ssl=True,
+ verify_certs=True,
+ connection_class=RequestsHttpConnection,
+ timeout=300,
+ )
+
@retry(
retry=retry_if_exception_type(TransportError),
stop=stop_after_attempt(30),
@@ -286,7 +311,7 @@ def create_vector_index(index_name: str, collection_id:
str, region: str):
reraise=True,
)
def _wait_for_index_readiness():
- oss_client.search(index=index_name, body={"query": {"match_all": {}},
"size": 0})
+ kb_role_oss_client.search(index=index_name, body={"query":
{"match_all": {}}, "size": 0})
_wait_for_index_readiness()
log.info("Index %s is ready.", index_name)
@@ -590,7 +615,12 @@ with DAG(
opensearch_policies,
collection,
await_collection,
- create_vector_index(index_name=index_name, collection_id=collection,
region=region_name),
+ create_vector_index(
+ index_name=index_name,
+ collection_id=collection,
+ region=region_name,
+ bedrock_role_arn=test_context[ROLE_ARN_KEY],
+ ),
copy_data_to_s3(bucket=bucket_name),
# TEST BODY
create_knowledge_base,
