diff --git a/providers/amazon/docs/operators/index.rst b/providers/amazon/docs/operators/index.rst index fa55e97d36f31..12fcb2dfaa581 100644 --- a/providers/amazon/docs/operators/index.rst +++ b/providers/amazon/docs/operators/index.rst @@ -29,4 +29,5 @@ Amazon AWS Operators emr/index redshift/index s3/index + s3tables/s3tables * diff --git a/providers/amazon/docs/operators/s3tables/s3tables.rst b/providers/amazon/docs/operators/s3tables/s3tables.rst new file mode 100644 index 0000000000000..48ac82e47a31e --- /dev/null +++ b/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] diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index 0ee5a70735e23..7fb3ee8d5df4a 100644 --- a/providers/amazon/provider.yaml +++ b/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 diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3tables.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3tables.py new file mode 100644 index 0000000000000..4c1bab617fde6 --- /dev/null +++ b/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 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 b09e13762686a..aac6b05f3c739 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/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"], diff --git a/providers/amazon/tests/system/amazon/aws/example_s3tables.py b/providers/amazon/tests/system/amazon/aws/example_s3tables.py new file mode 100644 index 0000000000000..d7a8526be8b8e --- /dev/null +++ b/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) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_s3tables.py b/providers/amazon/tests/unit/amazon/aws/operators/test_s3tables.py new file mode 100644 index 0000000000000..41f93515a0299 --- /dev/null +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_s3tables.py @@ -0,0 +1,81 @@ +# +# 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.base_aws import AwsBaseHook +from airflow.providers.amazon.aws.operators.s3tables import S3TablesCreateTableOperator + +from unit.amazon.aws.utils.test_template_fields import validate_template_fields + +TABLE_BUCKET_ARN = "arn:aws:s3tables:us-east-1:123456789012:bucket/test-bucket" +NAMESPACE = "test_namespace" +TABLE_NAME = "test_table" +TABLE_ARN = "arn:aws:s3tables:us-east-1:123456789012:bucket/test-bucket/table/test-id" + + +class TestS3TablesCreateTableOperator: + def setup_method(self): + self.operator = S3TablesCreateTableOperator( + task_id="test-create-table", + table_bucket_arn=TABLE_BUCKET_ARN, + namespace=NAMESPACE, + table_name=TABLE_NAME, + ) + + @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock) + def test_execute(self, mock_conn): + mock_client = mock.MagicMock() + mock_client.create_table.return_value = {"tableARN": TABLE_ARN, "versionToken": "v1"} + mock_conn.return_value = mock_client + + result = self.operator.execute({}) + mock_client.create_table.assert_called_once_with( + tableBucketARN=TABLE_BUCKET_ARN, + namespace=NAMESPACE, + name=TABLE_NAME, + format="ICEBERG", + ) + assert result == TABLE_ARN + + @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock) + def test_execute_with_metadata(self, mock_conn): + metadata = {"iceberg": {"schema": {"fields": [{"name": "id", "type": "int", "required": True}]}}} + op = S3TablesCreateTableOperator( + task_id="test-with-metadata", + table_bucket_arn=TABLE_BUCKET_ARN, + namespace=NAMESPACE, + table_name=TABLE_NAME, + metadata=metadata, + ) + mock_client = mock.MagicMock() + mock_client.create_table.return_value = {"tableARN": TABLE_ARN, "versionToken": "v1"} + mock_conn.return_value = mock_client + + op.execute({}) + mock_client.create_table.assert_called_once_with( + tableBucketARN=TABLE_BUCKET_ARN, + namespace=NAMESPACE, + name=TABLE_NAME, + format="ICEBERG", + metadata=metadata, + ) + + def test_template_fields(self): + validate_template_fields(self.operator)