Add S3TablesCreateTableOperator

providers/amazon/docs/operators/index.rst

@@ -29,4 +29,5 @@ Amazon AWS Operators
     emr/index
     redshift/index
     s3/index
+    s3tables/s3tables
     *

providers/amazon/docs/operators/s3tables/s3tables.rst

@@ -0,0 +1,34 @@
+ .. 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 Tables
+==================
+
+.. _howto/operator:S3TablesCreateTableOperator:
+
+Create an Amazon S3 Table
+=========================
+
+To create a new Iceberg table in an Amazon S3 Tables namespace you can use
+:class:`~airflow.providers.amazon.aws.operators.s3tables.S3TablesCreateTableOperator`.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3tables.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_s3tables_create_table]
+    :end-before: [END howto_operator_s3tables_create_table]

providers/amazon/provider.yaml

@@ -298,6 +298,12 @@ integrations:
     how-to-guide:
       - /docs/apache-airflow-providers-amazon/operators/s3/s3.rst
     tags: [aws]
+  - integration-name: Amazon S3 Tables
+    external-doc-url: https://aws.amazon.com/s3/features/tables/
+    logo: /docs/integration-logos/Amazon-Simple-Storage-Service-S3_light-bg@4x.png
+    how-to-guide:
+      - /docs/apache-airflow-providers-amazon/operators/s3tables/s3tables.rst
+    tags: [aws]
   - integration-name: Amazon Systems Manager (SSM)
     external-doc-url: https://aws.amazon.com/systems-manager/
     logo: /docs/integration-logos/AWS-Systems-Manager_light-bg@4x.png
@@ -450,6 +456,9 @@ operators:
   - integration-name: Amazon Simple Storage Service (S3)
     python-modules:
       - airflow.providers.amazon.aws.operators.s3
+  - integration-name: Amazon S3 Tables
+    python-modules:
+      - airflow.providers.amazon.aws.operators.s3tables
   - integration-name: Amazon SageMaker
     python-modules:
       - airflow.providers.amazon.aws.operators.sagemaker

providers/amazon/src/airflow/providers/amazon/aws/operators/s3tables.py

@@ -0,0 +1,104 @@
+#
+# 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 Tables operators."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.sdk import Context
+
+
+class S3TablesCreateTableOperator(AwsBaseOperator[AwsBaseHook]):
+    """
+    Create a new table in an Amazon S3 Tables namespace.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:S3TablesCreateTableOperator`
+
+    :param table_bucket_arn: The ARN of the table bucket to create the table in. (templated)
+    :param namespace: The namespace to associate with the table. (templated)
+    :param table_name: The name of the table. (templated)
+    :param format: The table format. (templated) Currently only ``ICEBERG`` is supported.
+    :param metadata: Optional Iceberg schema metadata. (templated)
+        Example: ``{"iceberg": {"schema": {"fields": [{"name": "id", "type": "int", "required": True}]}}}``
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param verify: Whether or not to verify SSL certificates. See:
+        https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    """
+
+    template_fields: Sequence[str] = aws_template_fields(
+        "table_bucket_arn", "namespace", "table_name", "format", "metadata"
+    )
+    template_fields_renderers = {"metadata": "json"}
+    aws_hook_class = AwsBaseHook
+
+    def __init__(
+        self,
+        *,
+        table_bucket_arn: str,
+        namespace: str,
+        table_name: str,
+        format: str = "ICEBERG",
+        metadata: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.table_bucket_arn = table_bucket_arn
+        self.namespace = namespace
+        self.table_name = table_name
+        self.format = format
+        self.metadata = metadata
+
+    @property
+    def _hook_parameters(self):
+        return {**super()._hook_parameters, "client_type": "s3tables"}
+
+    def execute(self, context: Context) -> str:
+        self.log.info(
+            "Creating S3 table %s in namespace %s (bucket %s)",
+            self.table_name,
+            self.namespace,
+            self.table_bucket_arn,
+        )
+        kwargs: dict[str, Any] = {
+            "tableBucketARN": self.table_bucket_arn,
+            "namespace": self.namespace,
+            "name": self.table_name,
+            "format": self.format,
+        }
+        if self.metadata:
+            kwargs["metadata"] = self.metadata
+        response = self.hook.conn.create_table(**kwargs)
+        table_arn = response["tableARN"]
+        self.log.info("Created table: %s", table_arn)
+        return table_arn

providers/amazon/src/airflow/providers/amazon/get_provider_info.py

@@ -247,6 +247,13 @@ def get_provider_info():
                 "how-to-guide": ["/docs/apache-airflow-providers-amazon/operators/s3/s3.rst"],
                 "tags": ["aws"],
             },
+            {
+                "integration-name": "Amazon S3 Tables",
+                "external-doc-url": "https://aws.amazon.com/s3/features/tables/",
+                "logo": "/docs/integration-logos/Amazon-Simple-Storage-Service-S3_light-bg@4x.png",
+                "how-to-guide": ["/docs/apache-airflow-providers-amazon/operators/s3tables/s3tables.rst"],
+                "tags": ["aws"],
+            },
             {
                 "integration-name": "Amazon Systems Manager (SSM)",
                 "external-doc-url": "https://aws.amazon.com/systems-manager/",
@@ -440,6 +447,10 @@ def get_provider_info():
                 "integration-name": "Amazon Simple Storage Service (S3)",
                 "python-modules": ["airflow.providers.amazon.aws.operators.s3"],
             },
+            {
+                "integration-name": "Amazon S3 Tables",
+                "python-modules": ["airflow.providers.amazon.aws.operators.s3tables"],
+            },
             {
                 "integration-name": "Amazon SageMaker",
                 "python-modules": ["airflow.providers.amazon.aws.operators.sagemaker"],

providers/amazon/tests/system/amazon/aws/example_s3tables.py

@@ -0,0 +1,142 @@
+# 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.s3tables import S3TablesCreateTableOperator
+from airflow.providers.common.compat.sdk import DAG, chain
+
+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]
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+
+DAG_ID = "example_s3tables"
+
+sys_test_context_task = SystemTestContextBuilder().build()
+
+SCHEMA = {
+    "iceberg": {
+        "schema": {
+            "fields": [
+                {"name": "id", "type": "int", "required": True},
+                {"name": "name", "type": "string", "required": False},
+            ]
+        }
+    }
+}
+
+with DAG(
+    dag_id=DAG_ID,
+    schedule="@once",
+    start_date=datetime(2021, 1, 1),
+    catchup=False,
+) as dag:
+    test_context = sys_test_context_task()
+    env_id = test_context[ENV_ID_KEY]
+
+    bucket_name = f"{env_id}-s3tables"
+    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."""
+        import boto3
+
+        boto3.client("s3tables").create_namespace(tableBucketARN=table_bucket_arn, namespace=[namespace])
+
+    @task(trigger_rule=TriggerRule.ALL_DONE)
+    def delete_table(table_bucket_arn: str, namespace: str, name: str):
+        """Delete the table."""
+        import boto3
+
+        client = boto3.client("s3tables")
+        try:
+            client.delete_table(tableBucketARN=table_bucket_arn, namespace=namespace, name=name)
+        except client.exceptions.NotFoundException:
+            pass
+
+    @task(trigger_rule=TriggerRule.ALL_DONE)
+    def delete_namespace(table_bucket_arn: str, namespace: str):
+        """Delete the namespace."""
+        import boto3
+
+        client = boto3.client("s3tables")
+        try:
+            client.delete_namespace(tableBucketARN=table_bucket_arn, namespace=namespace)
+        except client.exceptions.NotFoundException:
+            pass
+
+    @task(trigger_rule=TriggerRule.ALL_DONE)
+    def delete_table_bucket(table_bucket_arn: str):
+        """Delete the table bucket."""
+        import boto3
+
+        client = boto3.client("s3tables")
+        try:
+            client.delete_table_bucket(tableBucketARN=table_bucket_arn)
+        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]
+    create_table = S3TablesCreateTableOperator(
+        task_id="create_table",
+        table_bucket_arn=bucket_arn,
+        namespace=namespace,
+        table_name=table_name,
+        metadata=SCHEMA,
+    )
+    # [END howto_operator_s3tables_create_table]
+
+    chain(
+        # TEST SETUP
+        test_context,
+        bucket_arn,
+        setup_namespace,
+        # TEST BODY
+        create_table,
+        # TEST TEARDOWN
+        delete_table(table_bucket_arn=bucket_arn, namespace=namespace, name=table_name),
+        delete_namespace(table_bucket_arn=bucket_arn, namespace=namespace),
+        delete_table_bucket(table_bucket_arn=bucket_arn),
+    )
+
+    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)