diff --git a/providers/amazon/docs/operators/index.rst b/providers/amazon/docs/operators/index.rst index 12fcb2dfaa581..fa55e97d36f31 100644 --- a/providers/amazon/docs/operators/index.rst +++ b/providers/amazon/docs/operators/index.rst @@ -29,5 +29,4 @@ 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/s3_tables.rst similarity index 69% rename from providers/amazon/docs/operators/s3tables/s3tables.rst rename to providers/amazon/docs/operators/s3_tables.rst index 48ac82e47a31e..8fbaec8cad5b6 100644 --- a/providers/amazon/docs/operators/s3tables/s3tables.rst +++ b/providers/amazon/docs/operators/s3_tables.rst @@ -25,10 +25,24 @@ 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`. +:class:`~airflow.providers.amazon.aws.operators.s3_tables.S3TablesCreateTableOperator`. -.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3tables.py +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3_tables.py :language: python :dedent: 4 :start-after: [START howto_operator_s3tables_create_table] :end-before: [END howto_operator_s3tables_create_table] + +.. _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`. + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3_tables.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_s3tables_delete_table] + :end-before: [END howto_operator_s3tables_delete_table] diff --git a/providers/amazon/provider.yaml b/providers/amazon/provider.yaml index 1306418c31ef2..8169dc08a4af7 100644 --- a/providers/amazon/provider.yaml +++ b/providers/amazon/provider.yaml @@ -303,7 +303,7 @@ integrations: 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 + - /docs/apache-airflow-providers-amazon/operators/s3_tables.rst tags: [aws] - integration-name: Amazon Systems Manager (SSM) external-doc-url: https://aws.amazon.com/systems-manager/ @@ -461,7 +461,7 @@ operators: - airflow.providers.amazon.aws.operators.s3 - integration-name: Amazon S3 Tables python-modules: - - airflow.providers.amazon.aws.operators.s3tables + - airflow.providers.amazon.aws.operators.s3_tables - 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/s3_tables.py similarity index 68% rename from providers/amazon/src/airflow/providers/amazon/aws/operators/s3tables.py rename to providers/amazon/src/airflow/providers/amazon/aws/operators/s3_tables.py index 4c1bab617fde6..e6b5273c94ecf 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/s3tables.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/s3_tables.py @@ -25,6 +25,7 @@ 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 +from airflow.utils.helpers import prune_dict if TYPE_CHECKING: from airflow.sdk import Context @@ -102,3 +103,60 @@ def execute(self, context: Context) -> str: table_arn = response["tableARN"] self.log.info("Created table: %s", table_arn) return table_arn + + +class S3TablesDeleteTableOperator(AwsBaseOperator[AwsBaseHook]): + """ + Delete a table from an Amazon S3 Tables namespace. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:S3TablesDeleteTableOperator` + + :param table_bucket_arn: The ARN of the table bucket containing the table. (templated) + :param namespace: The namespace of the table. (templated) + :param table_name: The name of the table to delete. (templated) + :param version_token: Optional version token for optimistic concurrency. (templated) + """ + + template_fields: Sequence[str] = aws_template_fields( + "table_bucket_arn", "namespace", "table_name", "version_token" + ) + aws_hook_class = AwsBaseHook + + def __init__( + self, + *, + table_bucket_arn: str, + namespace: str, + table_name: str, + version_token: str | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.table_bucket_arn = table_bucket_arn + self.namespace = namespace + self.table_name = table_name + self.version_token = version_token + + @property + def _hook_parameters(self): + return {**super()._hook_parameters, "client_type": "s3tables"} + + def execute(self, context: Context) -> None: + self.log.info( + "Deleting S3 table %s from namespace %s (bucket %s)", + self.table_name, + self.namespace, + self.table_bucket_arn, + ) + kwargs: dict[str, Any] = prune_dict( + { + "tableBucketARN": self.table_bucket_arn, + "namespace": self.namespace, + "name": self.table_name, + "versionToken": self.version_token, + } + ) + self.hook.conn.delete_table(**kwargs) + self.log.info("Deleted table %s", self.table_name) 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 4b341628140d4..3468ef1d25a58 100644 --- a/providers/amazon/src/airflow/providers/amazon/get_provider_info.py +++ b/providers/amazon/src/airflow/providers/amazon/get_provider_info.py @@ -251,7 +251,7 @@ def get_provider_info(): "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"], + "how-to-guide": ["/docs/apache-airflow-providers-amazon/operators/s3_tables.rst"], "tags": ["aws"], }, { @@ -453,7 +453,7 @@ def get_provider_info(): }, { "integration-name": "Amazon S3 Tables", - "python-modules": ["airflow.providers.amazon.aws.operators.s3tables"], + "python-modules": ["airflow.providers.amazon.aws.operators.s3_tables"], }, { "integration-name": "Amazon SageMaker", diff --git a/providers/amazon/tests/system/amazon/aws/example_s3tables.py b/providers/amazon/tests/system/amazon/aws/example_s3_tables.py similarity index 87% rename from providers/amazon/tests/system/amazon/aws/example_s3tables.py rename to providers/amazon/tests/system/amazon/aws/example_s3_tables.py index d7a8526be8b8e..b91aff47ec35e 100644 --- a/providers/amazon/tests/system/amazon/aws/example_s3tables.py +++ b/providers/amazon/tests/system/amazon/aws/example_s3_tables.py @@ -18,7 +18,10 @@ from datetime import datetime -from airflow.providers.amazon.aws.operators.s3tables import S3TablesCreateTableOperator +from airflow.providers.amazon.aws.operators.s3_tables import ( + S3TablesCreateTableOperator, + S3TablesDeleteTableOperator, +) from airflow.providers.common.compat.sdk import DAG, chain from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS @@ -74,17 +77,6 @@ def create_namespace(table_bucket_arn: str, namespace: str): 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.""" @@ -120,6 +112,16 @@ def delete_table_bucket(table_bucket_arn: str): ) # [END howto_operator_s3tables_create_table] + # [START howto_operator_s3tables_delete_table] + delete_table = S3TablesDeleteTableOperator( + task_id="delete_table", + table_bucket_arn=bucket_arn, + namespace=namespace, + table_name=table_name, + trigger_rule=TriggerRule.ALL_DONE, + ) + # [END howto_operator_s3tables_delete_table] + chain( # TEST SETUP test_context, @@ -128,7 +130,7 @@ def delete_table_bucket(table_bucket_arn: str): # TEST BODY create_table, # TEST TEARDOWN - delete_table(table_bucket_arn=bucket_arn, namespace=namespace, name=table_name), + delete_table, delete_namespace(table_bucket_arn=bucket_arn, namespace=namespace), delete_table_bucket(table_bucket_arn=bucket_arn), ) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_s3tables.py b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py similarity index 65% rename from providers/amazon/tests/unit/amazon/aws/operators/test_s3tables.py rename to providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py index 41f93515a0299..5c6bb4b1acce4 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_s3tables.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py @@ -20,7 +20,10 @@ from unittest import mock from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook -from airflow.providers.amazon.aws.operators.s3tables import S3TablesCreateTableOperator +from airflow.providers.amazon.aws.operators.s3_tables import ( + S3TablesCreateTableOperator, + S3TablesDeleteTableOperator, +) from unit.amazon.aws.utils.test_template_fields import validate_template_fields @@ -79,3 +82,48 @@ def test_execute_with_metadata(self, mock_conn): def test_template_fields(self): validate_template_fields(self.operator) + + +class TestS3TablesDeleteTableOperator: + def setup_method(self): + self.operator = S3TablesDeleteTableOperator( + task_id="test-delete-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_conn.return_value = mock_client + + self.operator.execute({}) + mock_client.delete_table.assert_called_once_with( + tableBucketARN=TABLE_BUCKET_ARN, + namespace=NAMESPACE, + name=TABLE_NAME, + ) + + @mock.patch.object(AwsBaseHook, "conn", new_callable=mock.PropertyMock) + def test_execute_with_version_token(self, mock_conn): + op = S3TablesDeleteTableOperator( + task_id="test-delete-with-token", + table_bucket_arn=TABLE_BUCKET_ARN, + namespace=NAMESPACE, + table_name=TABLE_NAME, + version_token="v1", + ) + mock_client = mock.MagicMock() + mock_conn.return_value = mock_client + + op.execute({}) + mock_client.delete_table.assert_called_once_with( + tableBucketARN=TABLE_BUCKET_ARN, + namespace=NAMESPACE, + name=TABLE_NAME, + versionToken="v1", + ) + + def test_template_fields(self): + validate_template_fields(self.operator)