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 16450d2bdfa Add `S3TablesCreateTableBucketOperator` (#66119)
16450d2bdfa is described below
commit 16450d2bdfa23ced4b4e941cab729dc0a8692698
Author: John Jackson <[email protected]>
AuthorDate: Thu Apr 30 10:37:45 2026 -0700
Add `S3TablesCreateTableBucketOperator` (#66119)
---
providers/amazon/docs/operators/s3_tables.rst | 25 ++++++-
providers/amazon/provider.yaml | 3 +
.../providers/amazon/aws/hooks/s3_tables.py | 48 ++++++++++++
.../providers/amazon/aws/operators/s3_tables.py | 65 ++++++++++++++++-
.../airflow/providers/amazon/get_provider_info.py | 4 +
.../tests/system/amazon/aws/example_s3_tables.py | 28 ++++---
.../tests/unit/amazon/aws/hooks/test_s3_tables.py | 59 +++++++++++++++
.../unit/amazon/aws/operators/test_s3_tables.py | 85 ++++++++++++++++++++++
8 files changed, 298 insertions(+), 19 deletions(-)
diff --git a/providers/amazon/docs/operators/s3_tables.rst
b/providers/amazon/docs/operators/s3_tables.rst
index 8fbaec8cad5..ee4bc5e847d 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
<https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3tables.html>`__
diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml
index 151933dae65..8b563ae1b47 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 00000000000..72cfa710bed
--- /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") <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
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 e6b5273c94e..b38ef9f7a09 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 @@ class
S3TablesDeleteTableOperator(AwsBaseOperator[AwsBaseHook]):
)
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 ed1bc65269b..2fef1ae0d2e 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 6412cb06cad..c66f6959b65 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 __future__ import annotations
from datetime import datetime
from airflow.providers.amazon.aws.operators.s3_tables import (
+ S3TablesCreateTableBucketOperator,
S3TablesCreateTableOperator,
S3TablesDeleteTableOperator,
)
@@ -62,14 +63,6 @@ with DAG(
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 @@ with DAG(
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 @@ with DAG(
# [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 @@ with DAG(
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 00000000000..81407485068
--- /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 5c6bb4b1acc..6aa45edf4b2 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 __future__ import annotations
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 @@ class TestS3TablesDeleteTableOperator:
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)
