From 9e4fd4b434e80b67f7dcb98ce44dd66a9a47f084 Mon Sep 17 00:00:00 2001 From: john-jac <75442233+john-jac@users.noreply.github.com> Date: Tue, 28 Apr 2026 12:17:20 -0700 Subject: [PATCH] Add `S3VectorsDeleteVectorBucketOperator` - New operator to delete S3 Vectors vector buckets - System test updated to use operator instead of raw @task - Unit tests and docs included --- .../amazon/docs/operators/s3_vectors.rst | 14 ++++++++ .../amazon/aws/operators/s3_vectors.py | 35 +++++++++++++++++++ .../system/amazon/aws/example_s3_vectors.py | 23 ++++++------ .../amazon/aws/operators/test_s3_vectors.py | 27 ++++++++++++-- 4 files changed, 87 insertions(+), 12 deletions(-) diff --git a/providers/amazon/docs/operators/s3_vectors.rst b/providers/amazon/docs/operators/s3_vectors.rst index 177271ff7c91d..88db11ce77836 100644 --- a/providers/amazon/docs/operators/s3_vectors.rst +++ b/providers/amazon/docs/operators/s3_vectors.rst @@ -40,3 +40,17 @@ Reference ~~~~~~~~~ * `AWS boto3 Library Documentation for S3 Vectors `__ + +.. _howto/operator:S3VectorsDeleteVectorBucketOperator: + +Delete a Vector Bucket +~~~~~~~~~~~~~~~~~~~~~~ + +To delete an Amazon S3 Vectors vector bucket, use +:class:`~airflow.providers.amazon.aws.operators.s3_vectors.S3VectorsDeleteVectorBucketOperator`. + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3_vectors.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_s3vectors_delete_vector_bucket] + :end-before: [END howto_operator_s3vectors_delete_vector_bucket] diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py index 2c91c14642947..7708c56a58f61 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py @@ -91,3 +91,38 @@ def execute(self, context: Context) -> str: raise self.log.info("Vector bucket %s: %s", self.vector_bucket_name, vector_bucket_arn) return vector_bucket_arn + + +class S3VectorsDeleteVectorBucketOperator(AwsBaseOperator[AwsBaseHook]): + """ + Delete an Amazon S3 Vectors vector bucket. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:S3VectorsDeleteVectorBucketOperator` + + :param vector_bucket_name: The name of the vector bucket to delete. + """ + + aws_hook_class = AwsBaseHook + template_fields: tuple[str, ...] = ( + *AwsBaseOperator.template_fields, + "vector_bucket_name", + ) + + def __init__( + self, + *, + vector_bucket_name: str, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.vector_bucket_name = vector_bucket_name + + @property + def _hook_parameters(self) -> dict[str, Any]: + return {**super()._hook_parameters, "client_type": "s3vectors"} + + def execute(self, context: Context) -> None: + self.hook.conn.delete_vector_bucket(vectorBucketName=self.vector_bucket_name) + self.log.info("Deleted vector bucket %s", self.vector_bucket_name) diff --git a/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py b/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py index 48e3a423f4ce9..c2c75992f6cfe 100644 --- a/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py +++ b/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py @@ -18,16 +18,18 @@ from datetime import datetime -from airflow.providers.amazon.aws.operators.s3_vectors import S3VectorsCreateVectorBucketOperator +from airflow.providers.amazon.aws.operators.s3_vectors import ( + S3VectorsCreateVectorBucketOperator, + S3VectorsDeleteVectorBucketOperator, +) 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_s3vectors" @@ -51,17 +53,18 @@ ) # [END howto_operator_s3vectors_create_vector_bucket] - @task(trigger_rule=TriggerRule.ALL_DONE) - def delete_vector_bucket(name: str): - """Delete the vector bucket.""" - import boto3 - - boto3.client("s3vectors").delete_vector_bucket(vectorBucketName=name) + # [START howto_operator_s3vectors_delete_vector_bucket] + delete_vector_bucket = S3VectorsDeleteVectorBucketOperator( + task_id="delete_vector_bucket", + vector_bucket_name=bucket_name, + trigger_rule=TriggerRule.ALL_DONE, + ) + # [END howto_operator_s3vectors_delete_vector_bucket] chain( test_context, create_vector_bucket, - delete_vector_bucket(name=bucket_name), + delete_vector_bucket, ) from tests_common.test_utils.watcher import watcher diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py index 189fe9e0a89d4..f155bb1136794 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py @@ -20,7 +20,10 @@ from botocore.exceptions import ClientError -from airflow.providers.amazon.aws.operators.s3_vectors import S3VectorsCreateVectorBucketOperator +from airflow.providers.amazon.aws.operators.s3_vectors import ( + S3VectorsCreateVectorBucketOperator, + S3VectorsDeleteVectorBucketOperator, +) from unit.amazon.aws.utils.test_template_fields import validate_template_fields @@ -81,7 +84,6 @@ def test_execute_skip_existing(self): result = op.execute({}) assert result == BUCKET_ARN - mock_conn.get_vector_bucket.assert_called_once_with(vectorBucketName=BUCKET_NAME) def test_template_fields(self): op = S3VectorsCreateVectorBucketOperator( @@ -89,3 +91,24 @@ def test_template_fields(self): vector_bucket_name=BUCKET_NAME, ) validate_template_fields(op) + + +class TestS3VectorsDeleteVectorBucketOperator: + def test_execute(self): + op = S3VectorsDeleteVectorBucketOperator( + task_id="delete_vector_bucket", + vector_bucket_name=BUCKET_NAME, + ) + mock_conn = MagicMock() + op.hook.conn = mock_conn + + op.execute({}) + + mock_conn.delete_vector_bucket.assert_called_once_with(vectorBucketName=BUCKET_NAME) + + def test_template_fields(self): + op = S3VectorsDeleteVectorBucketOperator( + task_id="test", + vector_bucket_name=BUCKET_NAME, + ) + validate_template_fields(op)