This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 c711347b8e4 Add `BedrockDeleteGuardrailOperator` (#66140)
c711347b8e4 is described below
commit c711347b8e42adb3c5776111532f67fcc6cd6083
Author: John Jackson <[email protected]>
AuthorDate: Thu Apr 30 11:11:14 2026 -0700
Add `BedrockDeleteGuardrailOperator` (#66140)
---
providers/amazon/docs/operators/bedrock.rst | 14 ++++++++
.../providers/amazon/aws/operators/bedrock.py | 38 ++++++++++++++++++++
.../system/amazon/aws/example_bedrock_guardrail.py | 23 +++++++------
.../unit/amazon/aws/operators/test_bedrock.py | 40 ++++++++++++++++++++++
4 files changed, 105 insertions(+), 10 deletions(-)
diff --git a/providers/amazon/docs/operators/bedrock.rst
b/providers/amazon/docs/operators/bedrock.rst
index 374cf6b4a0d..f47dfb34ab1 100644
--- a/providers/amazon/docs/operators/bedrock.rst
+++ b/providers/amazon/docs/operators/bedrock.rst
@@ -353,6 +353,20 @@ To create an Amazon Bedrock guardrail, use
:start-after: [START howto_operator_bedrock_create_guardrail]
:end-before: [END howto_operator_bedrock_create_guardrail]
+.. _howto/operator:BedrockDeleteGuardrailOperator:
+
+Delete a Guardrail
+------------------
+
+To delete an Amazon Bedrock guardrail, use
+:class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockDeleteGuardrailOperator`.
+
+.. exampleinclude::
/../../amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_bedrock_delete_guardrail]
+ :end-before: [END howto_operator_bedrock_delete_guardrail]
+
Reference
---------
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 e561c4b1d6f..8b54d954142 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
@@ -1128,3 +1128,41 @@ class
BedrockCreateGuardrailOperator(AwsBaseOperator[BedrockHook]):
raise
self.log.info("Guardrail %s: %s", self.guardrail_name, guardrail_id)
return guardrail_id
+
+
+class BedrockDeleteGuardrailOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Delete an Amazon Bedrock guardrail.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:BedrockDeleteGuardrailOperator`
+
+ :param guardrail_identifier: The ID or ARN of the guardrail to delete.
(templated)
+ :param guardrail_version: Optional version number to delete a specific
version. (templated)
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] =
aws_template_fields("guardrail_identifier", "guardrail_version")
+
+ def __init__(
+ self,
+ *,
+ guardrail_identifier: str,
+ guardrail_version: str | None = None,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.guardrail_identifier = guardrail_identifier
+ self.guardrail_version = guardrail_version
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting Bedrock guardrail %s",
self.guardrail_identifier)
+ kwargs: dict[str, Any] = prune_dict(
+ {
+ "guardrailIdentifier": self.guardrail_identifier,
+ "guardrailVersion": self.guardrail_version,
+ }
+ )
+ self.hook.conn.delete_guardrail(**kwargs)
+ self.log.info("Deleted guardrail %s", self.guardrail_identifier)
diff --git
a/providers/amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
b/providers/amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
index 3473dce0f4f..d3db23c8812 100644
--- a/providers/amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
+++ b/providers/amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
@@ -18,16 +18,18 @@ from __future__ import annotations
from datetime import datetime
-from airflow.providers.amazon.aws.operators.bedrock import
BedrockCreateGuardrailOperator
+from airflow.providers.amazon.aws.operators.bedrock import (
+ BedrockCreateGuardrailOperator,
+ BedrockDeleteGuardrailOperator,
+)
from airflow.providers.common.compat.sdk import DAG, chain
from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
if AIRFLOW_V_3_0_PLUS:
- from airflow.sdk import TriggerRule, task
+ from airflow.sdk import TriggerRule
else:
- from airflow.decorators import task # type: ignore[attr-defined,no-redef]
from airflow.utils.trigger_rule import TriggerRule # type:
ignore[no-redef,attr-defined]
DAG_ID = "example_bedrock_guardrail"
@@ -58,17 +60,18 @@ with DAG(
)
# [END howto_operator_bedrock_create_guardrail]
- @task(trigger_rule=TriggerRule.ALL_DONE)
- def delete_guardrail(guardrail_id: str):
- """Delete the guardrail."""
- import boto3
-
-
boto3.client("bedrock").delete_guardrail(guardrailIdentifier=guardrail_id)
+ # [START howto_operator_bedrock_delete_guardrail]
+ delete_guardrail = BedrockDeleteGuardrailOperator(
+ task_id="delete_guardrail",
+ guardrail_identifier=create_guardrail.output,
+ trigger_rule=TriggerRule.ALL_DONE,
+ )
+ # [END howto_operator_bedrock_delete_guardrail]
chain(
test_context,
create_guardrail,
- delete_guardrail(guardrail_id=create_guardrail.output),
+ delete_guardrail,
)
from tests_common.test_utils.watcher import watcher
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
index 517ccf70dcf..4d41d8fbc69 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
@@ -34,6 +34,7 @@ from airflow.providers.amazon.aws.operators.bedrock import (
BedrockCreateKnowledgeBaseOperator,
BedrockCreateProvisionedModelThroughputOperator,
BedrockCustomizeModelOperator,
+ BedrockDeleteGuardrailOperator,
BedrockIngestDataOperator,
BedrockInvokeModelOperator,
BedrockRaGOperator,
@@ -873,3 +874,42 @@ class TestBedrockCreateGuardrailOperator:
def test_template_fields(self):
validate_template_fields(self.operator)
+
+
+GUARDRAIL_ID = "abc123"
+
+
+class TestBedrockDeleteGuardrailOperator:
+ def setup_method(self):
+ self.operator = BedrockDeleteGuardrailOperator(
+ task_id="delete_guardrail",
+ guardrail_identifier=GUARDRAIL_ID,
+ )
+
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute(self, mock_conn):
+ mock_client = mock.MagicMock()
+ mock_conn.return_value = mock_client
+
+ self.operator.execute({})
+
+
mock_client.delete_guardrail.assert_called_once_with(guardrailIdentifier=GUARDRAIL_ID)
+
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_with_version(self, mock_conn):
+ op = BedrockDeleteGuardrailOperator(
+ task_id="delete_guardrail",
+ guardrail_identifier=GUARDRAIL_ID,
+ guardrail_version="1",
+ )
+ mock_client = mock.MagicMock()
+ mock_conn.return_value = mock_client
+
+ op.execute({})
+
+ mock_client.delete_guardrail.assert_called_once_with(
+ guardrailIdentifier=GUARDRAIL_ID, guardrailVersion="1"
+ )
+
+ def test_template_fields(self):
+ validate_template_fields(self.operator)
