This is an automated email from the ASF dual-hosted git repository.
kaxil 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 f7f8ebff577 Add `S3VectorsCreateVectorBucketOperator` for Amazon S3
Vectors (#65968)
f7f8ebff577 is described below
commit f7f8ebff577c38d3058ffe0ef69e6c34cd8da280
Author: John Jackson <[email protected]>
AuthorDate: Tue Apr 28 10:31:44 2026 -0700
Add `S3VectorsCreateVectorBucketOperator` for Amazon S3 Vectors (#65968)
- New operator to create S3 Vectors vector buckets
- Idempotent: if_exists parameter (skip/fail) handles ConflictException
- Uses AwsBaseOperator[AwsBaseHook] with prune_dict
- Unit tests, system test with teardown, and docs included
- provider.yaml and get_provider_info.py updated
---
providers/amazon/docs/operators/s3_vectors.rst | 42 ++++++++++
providers/amazon/provider.yaml | 9 +++
.../providers/amazon/aws/operators/s3_vectors.py | 93 ++++++++++++++++++++++
.../airflow/providers/amazon/get_provider_info.py | 11 +++
.../tests/system/amazon/aws/example_s3_vectors.py | 73 +++++++++++++++++
.../unit/amazon/aws/operators/test_s3_vectors.py | 91 +++++++++++++++++++++
6 files changed, 319 insertions(+)
diff --git a/providers/amazon/docs/operators/s3_vectors.rst
b/providers/amazon/docs/operators/s3_vectors.rst
new file mode 100644
index 00000000000..177271ff7c9
--- /dev/null
+++ b/providers/amazon/docs/operators/s3_vectors.rst
@@ -0,0 +1,42 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+=================
+Amazon S3 Vectors
+=================
+
+Amazon S3 Vectors provides native vector storage in Amazon S3 for AI/ML
applications.
+Use the operators below to manage S3 Vectors resources.
+
+.. _howto/operator:S3VectorsCreateVectorBucketOperator:
+
+Create a Vector Bucket
+~~~~~~~~~~~~~~~~~~~~~~
+
+To create an Amazon S3 Vectors vector bucket, use
+:class:`~airflow.providers.amazon.aws.operators.s3_vectors.S3VectorsCreateVectorBucketOperator`.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3_vectors.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_s3vectors_create_vector_bucket]
+ :end-before: [END howto_operator_s3vectors_create_vector_bucket]
+
+Reference
+~~~~~~~~~
+
+* `AWS boto3 Library Documentation for S3 Vectors
<https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3vectors.html>`__
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index 1306418c31e..bb0e927f5ea 100644
--- a/providers/amazon/provider.yaml
+++ b/providers/amazon/provider.yaml
@@ -394,6 +394,12 @@ integrations:
how-to-guide:
- /docs/apache-airflow-providers-amazon/operators/mwaa.rst
tags: [aws]
+ - integration-name: Amazon S3 Vectors
+ external-doc-url:
https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-vectors.html
+ logo:
/docs/integration-logos/[email protected]
+ how-to-guide:
+ - /docs/apache-airflow-providers-amazon/operators/s3_vectors.rst
+ tags: [aws]
operators:
@@ -502,6 +508,9 @@ operators:
- integration-name: Amazon Neptune
python-modules:
- airflow.providers.amazon.aws.operators.neptune
+ - integration-name: Amazon S3 Vectors
+ python-modules:
+ - airflow.providers.amazon.aws.operators.s3_vectors
sensors:
- integration-name: Amazon Athena
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
new file mode 100644
index 00000000000..2c91c146429
--- /dev/null
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_vectors.py
@@ -0,0 +1,93 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Amazon S3 Vectors operators."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.utils.helpers import prune_dict
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class S3VectorsCreateVectorBucketOperator(AwsBaseOperator[AwsBaseHook]):
+ """
+ Create 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:S3VectorsCreateVectorBucketOperator`
+
+ :param vector_bucket_name: The name of the vector bucket to create (3-63
chars).
+ :param encryption_configuration: Optional encryption config dict with keys
+ ``sseType`` (``AES256`` or ``aws:kms``) and optionally ``kmsKeyArn``.
+ :param tags: Optional dict of tags to apply to the vector bucket.
+ :param if_exists: Behavior when the bucket already exists.
+ ``"fail"`` raises an error, ``"skip"`` returns the existing bucket ARN.
+ """
+
+ aws_hook_class = AwsBaseHook
+ template_fields: tuple[str, ...] = (
+ *AwsBaseOperator.template_fields,
+ "vector_bucket_name",
+ )
+
+ def __init__(
+ self,
+ *,
+ vector_bucket_name: str,
+ encryption_configuration: dict[str, Any] | None = None,
+ tags: dict[str, str] | None = None,
+ if_exists: Literal["fail", "skip"] = "skip",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.vector_bucket_name = vector_bucket_name
+ self.encryption_configuration = encryption_configuration
+ self.tags = tags
+ self.if_exists = if_exists
+
+ @property
+ def _hook_parameters(self) -> dict[str, Any]:
+ return {**super()._hook_parameters, "client_type": "s3vectors"}
+
+ def execute(self, context: Context) -> str:
+ kwargs: dict[str, Any] = prune_dict(
+ {
+ "vectorBucketName": self.vector_bucket_name,
+ "encryptionConfiguration": self.encryption_configuration,
+ "tags": self.tags,
+ }
+ )
+ try:
+ response = self.hook.conn.create_vector_bucket(**kwargs)
+ vector_bucket_arn = response["vectorBucketArn"]
+ except ClientError as e:
+ if e.response["Error"]["Code"] == "ConflictException" and
self.if_exists == "skip":
+ self.log.info("Vector bucket %s already exists, skipping.",
self.vector_bucket_name)
+ response =
self.hook.conn.get_vector_bucket(vectorBucketName=self.vector_bucket_name)
+ vector_bucket_arn = response["vectorBucketArn"]
+ else:
+ raise
+ self.log.info("Vector bucket %s: %s", self.vector_bucket_name,
vector_bucket_arn)
+ return vector_bucket_arn
diff --git a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
index 4b341628140..a3a0a2d9a9d 100644
--- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
+++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py
@@ -365,6 +365,13 @@ def get_provider_info():
"how-to-guide":
["/docs/apache-airflow-providers-amazon/operators/mwaa.rst"],
"tags": ["aws"],
},
+ {
+ "integration-name": "Amazon S3 Vectors",
+ "external-doc-url":
"https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-vectors.html";,
+ "logo":
"/docs/integration-logos/[email protected]",
+ "how-to-guide":
["/docs/apache-airflow-providers-amazon/operators/s3_vectors.rst"],
+ "tags": ["aws"],
+ },
],
"operators": [
{
@@ -510,6 +517,10 @@ def get_provider_info():
"integration-name": "Amazon Neptune",
"python-modules":
["airflow.providers.amazon.aws.operators.neptune"],
},
+ {
+ "integration-name": "Amazon S3 Vectors",
+ "python-modules":
["airflow.providers.amazon.aws.operators.s3_vectors"],
+ },
],
"sensors": [
{
diff --git a/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py
b/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py
new file mode 100644
index 00000000000..48e3a423f4c
--- /dev/null
+++ b/providers/amazon/tests/system/amazon/aws/example_s3_vectors.py
@@ -0,0 +1,73 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.s3_vectors import
S3VectorsCreateVectorBucketOperator
+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
+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"
+
+sys_test_context_task = SystemTestContextBuilder().build()
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule=None,
+ start_date=datetime(2024, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+ env_id = test_context[ENV_ID_KEY]
+ bucket_name = f"{env_id}-test-vectors"
+
+ # [START howto_operator_s3vectors_create_vector_bucket]
+ create_vector_bucket = S3VectorsCreateVectorBucketOperator(
+ task_id="create_vector_bucket",
+ vector_bucket_name=bucket_name,
+ )
+ # [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)
+
+ chain(
+ test_context,
+ create_vector_bucket,
+ delete_vector_bucket(name=bucket_name),
+ )
+
+ from tests_common.test_utils.watcher import watcher
+
+ list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+test_run = get_test_run(dag)
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
new file mode 100644
index 00000000000..189fe9e0a89
--- /dev/null
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_vectors.py
@@ -0,0 +1,91 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.operators.s3_vectors import
S3VectorsCreateVectorBucketOperator
+
+from unit.amazon.aws.utils.test_template_fields import validate_template_fields
+
+BUCKET_NAME = "test-vector-bucket"
+BUCKET_ARN =
"arn:aws:s3vectors:us-east-1:123456789012:bucket/test-vector-bucket"
+
+
+class TestS3VectorsCreateVectorBucketOperator:
+ def test_execute(self):
+ op = S3VectorsCreateVectorBucketOperator(
+ task_id="create_vector_bucket",
+ vector_bucket_name=BUCKET_NAME,
+ )
+ mock_conn = MagicMock()
+ mock_conn.create_vector_bucket.return_value = {"vectorBucketArn":
BUCKET_ARN}
+ op.hook.conn = mock_conn
+
+ result = op.execute({})
+
+
mock_conn.create_vector_bucket.assert_called_once_with(vectorBucketName=BUCKET_NAME)
+ assert result == BUCKET_ARN
+
+ def test_execute_with_all_params(self):
+ encryption = {"sseType": "aws:kms", "kmsKeyArn":
"arn:aws:kms:us-east-1:123456789012:key/abc"}
+ op = S3VectorsCreateVectorBucketOperator(
+ task_id="create_vector_bucket",
+ vector_bucket_name=BUCKET_NAME,
+ encryption_configuration=encryption,
+ tags={"env": "test"},
+ )
+ mock_conn = MagicMock()
+ mock_conn.create_vector_bucket.return_value = {"vectorBucketArn":
BUCKET_ARN}
+ op.hook.conn = mock_conn
+
+ result = op.execute({})
+
+ mock_conn.create_vector_bucket.assert_called_once_with(
+ vectorBucketName=BUCKET_NAME,
+ encryptionConfiguration=encryption,
+ tags={"env": "test"},
+ )
+ assert result == BUCKET_ARN
+
+ def test_execute_skip_existing(self):
+ op = S3VectorsCreateVectorBucketOperator(
+ task_id="create_vector_bucket",
+ vector_bucket_name=BUCKET_NAME,
+ if_exists="skip",
+ )
+ mock_conn = MagicMock()
+ mock_conn.create_vector_bucket.side_effect = ClientError(
+ {"Error": {"Code": "ConflictException", "Message": "Already
exists"}},
+ "CreateVectorBucket",
+ )
+ mock_conn.get_vector_bucket.return_value = {"vectorBucketArn":
BUCKET_ARN}
+ op.hook.conn = mock_conn
+
+ 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(
+ task_id="test",
+ vector_bucket_name=BUCKET_NAME,
+ )
+ validate_template_fields(op)
