From 189488fcb1d0b6a7f4a774112e1d41f029c923ad Mon Sep 17 00:00:00 2001 From: john-jac <75442233+john-jac@users.noreply.github.com> Date: Wed, 29 Apr 2026 17:45:10 -0700 Subject: [PATCH] Add `S3TablesCreateTableBucketOperator` --- providers/amazon/docs/operators/s3_tables.rst | 25 +++++- providers/amazon/provider.yaml | 3 + .../providers/amazon/aws/hooks/s3_tables.py | 48 +++++++++++ .../amazon/aws/operators/s3_tables.py | 65 +++++++++++++- .../providers/amazon/get_provider_info.py | 4 + .../system/amazon/aws/example_s3_tables.py | 30 +++---- .../unit/amazon/aws/hooks/test_s3_tables.py | 59 +++++++++++++ .../amazon/aws/operators/test_s3_tables.py | 85 +++++++++++++++++++ 8 files changed, 299 insertions(+), 20 deletions(-) create mode 100644 providers/amazon/src/airflow/providers/amazon/aws/hooks/s3_tables.py create mode 100644 providers/amazon/tests/unit/amazon/aws/hooks/test_s3_tables.py diff --git a/providers/amazon/docs/operators/s3_tables.rst b/providers/amazon/docs/operators/s3_tables.rst index 8fbaec8cad5b6..ee4bc5e847d90 100644 --- a/providers/amazon/docs/operators/s3_tables.rst +++ b/providers/amazon/docs/operators/s3_tables.rst @@ -19,10 +19,24 @@ Amazon S3 Tables ================== +.. _howto/operator:S3TablesCreateTableBucketOperator: + +Create a Table Bucket +--------------------- + +To create an Amazon S3 Tables table bucket, use +:class:`~airflow.providers.amazon.aws.operators.s3_tables.S3TablesCreateTableBucketOperator`. + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3_tables.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_s3tables_create_table_bucket] + :end-before: [END howto_operator_s3tables_create_table_bucket] + .. _howto/operator:S3TablesCreateTableOperator: -Create an Amazon S3 Table -========================= +Create a Table +-------------- To create a new Iceberg table in an Amazon S3 Tables namespace you can use :class:`~airflow.providers.amazon.aws.operators.s3_tables.S3TablesCreateTableOperator`. @@ -36,7 +50,7 @@ To create a new Iceberg table in an Amazon S3 Tables namespace you can use .. _howto/operator:S3TablesDeleteTableOperator: Delete a Table -~~~~~~~~~~~~~~ +-------------- To delete a table from an Amazon S3 Tables namespace, use :class:`~airflow.providers.amazon.aws.operators.s3_tables.S3TablesDeleteTableOperator`. @@ -46,3 +60,8 @@ To delete a table from an Amazon S3 Tables namespace, use :dedent: 4 :start-after: [START howto_operator_s3tables_delete_table] :end-before: [END howto_operator_s3tables_delete_table] + +Reference +--------- + +* `AWS boto3 Library Documentation for S3 Tables `__ diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index 151933dae65a4..8b563ae1b47dd 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -714,6 +714,9 @@ hooks: - integration-name: Amazon Simple Storage Service (S3) python-modules: - airflow.providers.amazon.aws.hooks.s3 + - integration-name: Amazon S3 Tables + python-modules: + - airflow.providers.amazon.aws.hooks.s3_tables - integration-name: Amazon SageMaker python-modules: - airflow.providers.amazon.aws.hooks.sagemaker diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3_tables.py b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3_tables.py new file mode 100644 index 0000000000000..72cfa710bed58 --- /dev/null +++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/s3_tables.py @@ -0,0 +1,48 @@ +# +# 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 airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook + + +class S3TablesHook(AwsBaseHook): + """ + Interact with Amazon S3 Tables. + + Provide thin wrapper around + :external+boto3:py:class:`boto3.client("s3tables") `. + + Additional arguments (such as ``aws_conn_id``) may be specified and + are passed down to the underlying AwsBaseHook. + + .. seealso:: + - :class:`airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook` + """ + + def __init__(self, *args, **kwargs) -> None: + kwargs["client_type"] = "s3tables" + super().__init__(*args, **kwargs) + + def get_table_bucket_arn_by_name(self, table_bucket_name: str) -> str | None: + """Get the table bucket ARN by name, or None if not found.""" + paginator = self.conn.get_paginator("list_table_buckets") + for page in paginator.paginate(): + for tb in page.get("tableBuckets", []): + if tb["name"] == table_bucket_name: + return tb["arn"] + return None diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_tables.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_tables.py index e6b5273c94ecf..b38ef9f7a093d 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_tables.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_tables.py @@ -20,9 +20,12 @@ from __future__ import annotations from collections.abc import Sequence -from typing import TYPE_CHECKING, Any +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.hooks.s3_tables import S3TablesHook from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator from airflow.providers.amazon.aws.utils.mixins import aws_template_fields from airflow.utils.helpers import prune_dict @@ -160,3 +163,63 @@ def execute(self, context: Context) -> None: ) self.hook.conn.delete_table(**kwargs) self.log.info("Deleted table %s", self.table_name) + + +class S3TablesCreateTableBucketOperator(AwsBaseOperator[S3TablesHook]): + """ + Create an Amazon S3 Tables table bucket. + + A table bucket is the top-level container for S3 Tables namespaces and tables. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:S3TablesCreateTableBucketOperator` + + :param table_bucket_name: The name of the table bucket. (templated) + :param encryption_configuration: Optional encryption configuration dict with + ``sseAlgorithm`` and optional ``kmsKeyArn``. (templated) + :param if_exists: Behavior when a table bucket with the same name already exists. + ``"fail"`` raises an error, ``"skip"`` returns the existing bucket ARN. + """ + + template_fields: Sequence[str] = aws_template_fields("table_bucket_name") + template_fields_renderers = {"encryption_configuration": "json"} + aws_hook_class = S3TablesHook + + def __init__( + self, + *, + table_bucket_name: str, + encryption_configuration: dict[str, str] | None = None, + tags: dict[str, str] | None = None, + if_exists: Literal["fail", "skip"] = "skip", + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.table_bucket_name = table_bucket_name + self.encryption_configuration = encryption_configuration + self.tags = tags + self.if_exists = if_exists + + def execute(self, context: Context) -> str: + self.log.info("Creating S3 Tables table bucket %s", self.table_bucket_name) + kwargs: dict[str, Any] = prune_dict( + { + "name": self.table_bucket_name, + "encryptionConfiguration": self.encryption_configuration, + "tags": self.tags, + } + ) + try: + response = self.hook.conn.create_table_bucket(**kwargs) + bucket_arn = response["arn"] + except ClientError as e: + if e.response["Error"]["Code"] == "ConflictException" and self.if_exists == "skip": + self.log.info("Table bucket %s already exists, skipping.", self.table_bucket_name) + bucket_arn = self.hook.get_table_bucket_arn_by_name(self.table_bucket_name) + if bucket_arn is None: + raise + else: + raise + self.log.info("Table bucket %s: %s", self.table_bucket_name, bucket_arn) + return 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 ed1bc65269b43..2fef1ae0d2e55 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py @@ -780,6 +780,10 @@ def get_provider_info(): "integration-name": "Amazon Simple Storage Service (S3)", "python-modules": ["airflow.providers.amazon.aws.hooks.s3"], }, + { + "integration-name": "Amazon S3 Tables", + "python-modules": ["airflow.providers.amazon.aws.hooks.s3_tables"], + }, { "integration-name": "Amazon SageMaker", "python-modules": ["airflow.providers.amazon.aws.hooks.sagemaker"], diff --git a/providers/amazon/tests/system/amazon/aws/example_s3_tables.py b/providers/amazon/tests/system/amazon/aws/example_s3_tables.py index b91aff47ec35e..c66f6959b65f1 100644 --- a/providers/amazon/tests/system/amazon/aws/example_s3_tables.py +++ b/providers/amazon/tests/system/amazon/aws/example_s3_tables.py @@ -19,6 +19,7 @@ from datetime import datetime from airflow.providers.amazon.aws.operators.s3_tables import ( + S3TablesCreateTableBucketOperator, S3TablesCreateTableOperator, S3TablesDeleteTableOperator, ) @@ -34,7 +35,7 @@ from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder -DAG_ID = "example_s3tables" +DAG_ID = "example_s3_tables" sys_test_context_task = SystemTestContextBuilder().build() @@ -62,14 +63,6 @@ namespace = f"{env_id}_ns" table_name = f"{env_id}_tbl" - @task - def create_table_bucket(name: str) -> str: - """Create an S3 Tables bucket and return its ARN.""" - import boto3 - - client = boto3.client("s3tables") - return client.create_table_bucket(name=name)["arn"] - @task def create_namespace(table_bucket_arn: str, namespace: str): """Create a namespace in the table bucket.""" @@ -99,13 +92,18 @@ def delete_table_bucket(table_bucket_arn: str): except client.exceptions.NotFoundException: pass - bucket_arn = create_table_bucket(name=bucket_name) - setup_namespace = create_namespace(table_bucket_arn=bucket_arn, namespace=namespace) + # [START howto_operator_s3tables_create_table_bucket] + create_table_bucket = S3TablesCreateTableBucketOperator( + task_id="create_table_bucket", + table_bucket_name=bucket_name, + ) + # [END howto_operator_s3tables_create_table_bucket] + setup_namespace = create_namespace(table_bucket_arn=create_table_bucket.output, namespace=namespace) # [START howto_operator_s3tables_create_table] create_table = S3TablesCreateTableOperator( task_id="create_table", - table_bucket_arn=bucket_arn, + table_bucket_arn=create_table_bucket.output, namespace=namespace, table_name=table_name, metadata=SCHEMA, @@ -115,7 +113,7 @@ def delete_table_bucket(table_bucket_arn: str): # [START howto_operator_s3tables_delete_table] delete_table = S3TablesDeleteTableOperator( task_id="delete_table", - table_bucket_arn=bucket_arn, + table_bucket_arn=create_table_bucket.output, namespace=namespace, table_name=table_name, trigger_rule=TriggerRule.ALL_DONE, @@ -125,14 +123,14 @@ def delete_table_bucket(table_bucket_arn: str): chain( # TEST SETUP test_context, - bucket_arn, + create_table_bucket, setup_namespace, # TEST BODY create_table, # TEST TEARDOWN delete_table, - delete_namespace(table_bucket_arn=bucket_arn, namespace=namespace), - delete_table_bucket(table_bucket_arn=bucket_arn), + delete_namespace(table_bucket_arn=create_table_bucket.output, namespace=namespace), + delete_table_bucket(table_bucket_arn=create_table_bucket.output), ) from tests_common.test_utils.watcher import watcher diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_s3_tables.py b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3_tables.py new file mode 100644 index 0000000000000..81407485068cf --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_s3_tables.py @@ -0,0 +1,59 @@ +# +# 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 import mock + +from airflow.providers.amazon.aws.hooks.s3_tables import S3TablesHook + + +class TestS3TablesHook: + def test_client_type(self): + hook = S3TablesHook() + assert hook.client_type == "s3tables" + + @mock.patch.object(S3TablesHook, "conn", new_callable=mock.PropertyMock) + def test_get_table_bucket_arn_by_name_found(self, mock_conn): + mock_client = mock.MagicMock() + mock_paginator = mock.MagicMock() + mock_paginator.paginate.return_value = [ + { + "tableBuckets": [ + {"name": "my-bucket", "arn": "arn:aws:s3tables:us-east-1:123:bucket/my-bucket"} + ] + } + ] + mock_client.get_paginator.return_value = mock_paginator + mock_conn.return_value = mock_client + + hook = S3TablesHook() + assert ( + hook.get_table_bucket_arn_by_name("my-bucket") + == "arn:aws:s3tables:us-east-1:123:bucket/my-bucket" + ) + + @mock.patch.object(S3TablesHook, "conn", new_callable=mock.PropertyMock) + def test_get_table_bucket_arn_by_name_not_found(self, mock_conn): + mock_client = mock.MagicMock() + mock_paginator = mock.MagicMock() + mock_paginator.paginate.return_value = [{"tableBuckets": []}] + mock_client.get_paginator.return_value = mock_paginator + mock_conn.return_value = mock_client + + hook = S3TablesHook() + assert hook.get_table_bucket_arn_by_name("nonexistent") is None diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py index 5c6bb4b1acce4..6aa45edf4b2f7 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py @@ -19,8 +19,13 @@ from unittest import mock +import pytest +from botocore.exceptions import ClientError + from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook +from airflow.providers.amazon.aws.hooks.s3_tables import S3TablesHook from airflow.providers.amazon.aws.operators.s3_tables import ( + S3TablesCreateTableBucketOperator, S3TablesCreateTableOperator, S3TablesDeleteTableOperator, ) @@ -127,3 +132,83 @@ def test_execute_with_version_token(self, mock_conn): def test_template_fields(self): validate_template_fields(self.operator) + + +BUCKET_NAME = "test-table-bucket" +BUCKET_ARN = "arn:aws:s3tables:us-east-1:123456789012:bucket/test-table-bucket" + + +class TestS3TablesCreateTableBucketOperator: + def setup_method(self): + self.operator = S3TablesCreateTableBucketOperator( + task_id="create_table_bucket", + table_bucket_name=BUCKET_NAME, + ) + + @mock.patch.object(S3TablesHook, "conn", new_callable=mock.PropertyMock) + def test_execute(self, mock_conn): + mock_client = mock.MagicMock() + mock_client.create_table_bucket.return_value = {"arn": BUCKET_ARN} + mock_conn.return_value = mock_client + + result = self.operator.execute({}) + + mock_client.create_table_bucket.assert_called_once_with(name=BUCKET_NAME) + assert result == BUCKET_ARN + + @mock.patch.object(S3TablesHook, "conn", new_callable=mock.PropertyMock) + def test_execute_with_encryption(self, mock_conn): + enc = {"sseAlgorithm": "aws:kms", "kmsKeyArn": "arn:aws:kms:us-east-1:123:key/abc"} + op = S3TablesCreateTableBucketOperator( + task_id="create_table_bucket", + table_bucket_name=BUCKET_NAME, + encryption_configuration=enc, + tags={"env": "test"}, + ) + mock_client = mock.MagicMock() + mock_client.create_table_bucket.return_value = {"arn": BUCKET_ARN} + mock_conn.return_value = mock_client + + result = op.execute({}) + + mock_client.create_table_bucket.assert_called_once_with( + name=BUCKET_NAME, + encryptionConfiguration=enc, + tags={"env": "test"}, + ) + assert result == BUCKET_ARN + + @mock.patch.object(S3TablesHook, "get_table_bucket_arn_by_name", return_value=BUCKET_ARN) + @mock.patch.object(S3TablesHook, "conn", new_callable=mock.PropertyMock) + def test_execute_skip_existing(self, mock_conn, mock_get_arn): + mock_client = mock.MagicMock() + mock_client.create_table_bucket.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "Already exists"}}, + "CreateTableBucket", + ) + mock_conn.return_value = mock_client + + result = self.operator.execute({}) + + assert result == BUCKET_ARN + mock_get_arn.assert_called_once_with(BUCKET_NAME) + + @mock.patch.object(S3TablesHook, "conn", new_callable=mock.PropertyMock) + def test_execute_fail_on_conflict(self, mock_conn): + op = S3TablesCreateTableBucketOperator( + task_id="create_table_bucket", + table_bucket_name=BUCKET_NAME, + if_exists="fail", + ) + mock_client = mock.MagicMock() + mock_client.create_table_bucket.side_effect = ClientError( + {"Error": {"Code": "ConflictException", "Message": "Already exists"}}, + "CreateTableBucket", + ) + mock_conn.return_value = mock_client + + with pytest.raises(ClientError): + op.execute({}) + + def test_template_fields(self): + validate_template_fields(self.operator)