Add S3TablesCreateTableBucketOperator

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 <https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3tables.html>`__

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

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") <S3Tables.Client>`.
+
+    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

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

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"],

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

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