Add S3TablesPutTableBucketPolicyOperator

providers/amazon/docs/operators/s3_tables.rst

@@ -117,6 +117,21 @@ To rename a table in an Amazon S3 Tables namespace, use
     :start-after: [START howto_operator_s3tables_rename_table]
     :end-before: [END howto_operator_s3tables_rename_table]
 
+
+.. _howto/operator:S3TablesPutTableBucketPolicyOperator:
+
+Put a Table Bucket Policy
+-------------------------
+
+To set a resource policy on an Amazon S3 Tables table bucket, use
+:class:`~airflow.providers.amazon.aws.operators.s3_tables.S3TablesPutTableBucketPolicyOperator`.
+
+.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_s3_tables.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_s3tables_put_table_bucket_policy]
+    :end-before: [END howto_operator_s3tables_put_table_bucket_policy]
+
 Reference
 ---------
 

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

@@ -383,3 +383,39 @@ def execute(self, context: Context) -> None:
         )
         self.hook.conn.rename_table(**kwargs)
         self.log.info("Renamed table %s to %s", self.table_name, self.new_name)
+
+
+class S3TablesPutTableBucketPolicyOperator(AwsBaseOperator[S3TablesHook]):
+    """
+    Set a resource policy on an Amazon S3 Tables table bucket.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:S3TablesPutTableBucketPolicyOperator`
+
+    :param table_bucket_arn: The ARN of the table bucket. (templated)
+    :param resource_policy: The JSON resource policy string. (templated)
+    """
+
+    aws_hook_class = S3TablesHook
+    template_fields: Sequence[str] = aws_template_fields("table_bucket_arn", "resource_policy")
+    template_fields_renderers = {"resource_policy": "json"}
+
+    def __init__(
+        self,
+        *,
+        table_bucket_arn: str,
+        resource_policy: str,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.table_bucket_arn = table_bucket_arn
+        self.resource_policy = resource_policy
+
+    def execute(self, context: Context) -> None:
+        self.log.info("Setting policy on table bucket %s", self.table_bucket_arn)
+        self.hook.conn.put_table_bucket_policy(
+            tableBucketARN=self.table_bucket_arn,
+            resourcePolicy=self.resource_policy,
+        )
+        self.log.info("Policy set on table bucket %s", self.table_bucket_arn)

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

@@ -25,6 +25,7 @@
     S3TablesDeleteNamespaceOperator,
     S3TablesDeleteTableBucketOperator,
     S3TablesDeleteTableOperator,
+    S3TablesPutTableBucketPolicyOperator,
     S3TablesRenameTableOperator,
 )
 from airflow.providers.common.compat.sdk import DAG, chain
@@ -81,6 +82,15 @@ def create_namespace(table_bucket_arn: str, namespace: str):
         table_bucket_name=bucket_name,
     )
     # [END howto_operator_s3tables_create_table_bucket]
+
+    # [START howto_operator_s3tables_put_table_bucket_policy]
+    put_policy = S3TablesPutTableBucketPolicyOperator(
+        task_id="put_table_bucket_policy",
+        table_bucket_arn=create_table_bucket.output,
+        resource_policy='{"Version":"2012-10-17","Statement":[]}',
+    )
+    # [END howto_operator_s3tables_put_table_bucket_policy]
+
     # [START howto_operator_s3tables_create_namespace]
     setup_namespace = S3TablesCreateNamespaceOperator(
         task_id="create_namespace",
@@ -140,6 +150,7 @@ def create_namespace(table_bucket_arn: str, namespace: str):
         # TEST SETUP
         test_context,
         create_table_bucket,
+        put_policy,
         setup_namespace,
         # TEST BODY
         create_table,

providers/amazon/tests/unit/amazon/aws/operators/test_s3_tables.py

@@ -31,6 +31,7 @@
     S3TablesDeleteNamespaceOperator,
     S3TablesDeleteTableBucketOperator,
     S3TablesDeleteTableOperator,
+    S3TablesPutTableBucketPolicyOperator,
     S3TablesRenameTableOperator,
 )
 
@@ -392,3 +393,30 @@ def test_execute_with_optional_args(self, mock_conn):
 
     def test_template_fields(self):
         validate_template_fields(self.operator)
+
+
+POLICY = '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":"*","Action":"s3tables:*","Resource":"*"}]}'
+
+
+class TestS3TablesPutTableBucketPolicyOperator:
+    def setup_method(self):
+        self.operator = S3TablesPutTableBucketPolicyOperator(
+            task_id="put_policy",
+            table_bucket_arn=TABLE_BUCKET_ARN,
+            resource_policy=POLICY,
+        )
+
+    @mock.patch.object(S3TablesHook, "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.put_table_bucket_policy.assert_called_once_with(
+            tableBucketARN=TABLE_BUCKET_ARN,
+            resourcePolicy=POLICY,
+        )
+
+    def test_template_fields(self):
+        validate_template_fields(self.operator)