Add GCSBucketAddIamBindingOperator to apply IAM roles for gcs buckets

providers/google/docs/operators/cloud/cloud_memorystore.rst

@@ -51,22 +51,22 @@ Here is an example of instance
 Configuration of bucket permissions for import / export
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-It is necessary to configure permissions for the bucket to import and export data. Too find the service
+It is necessary to configure permissions for the bucket to import and export data. To find the service
 account for your instance, run the :class:`~airflow.providers.google.cloud.operators.cloud_memorystore.CloudMemorystoreCreateInstanceOperator` or
 :class:`~airflow.providers.google.cloud.operators.cloud_memorystore.CloudMemorystoreGetInstanceOperator` and
 make a use of the service account listed under ``persistenceIamIdentity``.
 
-You can use :class:`~airflow.providers.google.cloud.operators.gcs.GCSBucketCreateAclEntryOperator`
-operator to set permissions.
+Grant the service account bucket-level IAM roles that provide access to the bucket metadata and objects.
+This works with uniform bucket-level access, where bucket and object ACLs are disabled.
 
 .. exampleinclude:: /../../google/tests/system/google/cloud/cloud_memorystore/example_cloud_memorystore_redis.py
     :language: python
     :dedent: 4
-    :start-after: [START howto_operator_set_acl_permission]
-    :end-before: [END howto_operator_set_acl_permission]
+    :start-after: [START howto_operator_set_iam_permissions]
+    :end-before: [END howto_operator_set_iam_permissions]
 
 For further information look at: `Granting restricted permissions for import and export
-<https://cloud.google.com/memorystore/docs/redis/import-export-restricted-permissions>`__
+<https://cloud.google.com/memorystore/docs/redis/access-control#required_permissions_for_import_and_export>`__
 
 .. _howto/operator:CloudMemorystoreCreateInstanceOperator:
 

providers/google/docs/operators/cloud/cloud_sql.rst

@@ -264,9 +264,9 @@ of the Cloud SQL instance is authorized to write to the selected GCS bucket.
 It is not the service account configured in Airflow that communicates with GCS,
 but rather the service account of the particular Cloud SQL instance.
 
-To grant the service account with the appropriate WRITE permissions for the GCS bucket
-you can use the :class:`~airflow.providers.google.cloud.operators.gcs.GCSBucketCreateAclEntryOperator`,
-as shown in the example:
+Grant the Cloud SQL instance service account the ``roles/storage.objectAdmin`` IAM role
+on the GCS bucket. This works with uniform bucket-level access, where bucket and object
+ACLs are disabled:
 
 .. exampleinclude:: /../../google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py
     :language: python
@@ -349,15 +349,14 @@ of the Cloud SQL instance is authorized to read from the selected GCS object.
 It is not the service account configured in Airflow that communicates with GCS,
 but rather the service account of the particular Cloud SQL instance.
 
-To grant the service account with the appropriate READ permissions for the GCS object
-you can use the :class:`~airflow.providers.google.cloud.operators.gcs.GCSBucketCreateAclEntryOperator`,
-as shown in the example:
+The ``roles/storage.objectAdmin`` IAM role granted for export also provides the object read
+permissions required for import:
 
 .. exampleinclude:: /../../google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py
     :language: python
     :dedent: 4
-    :start-after: [START howto_operator_cloudsql_import_gcs_permissions]
-    :end-before: [END howto_operator_cloudsql_import_gcs_permissions]
+    :start-after: [START howto_operator_cloudsql_export_gcs_permissions]
+    :end-before: [END howto_operator_cloudsql_export_gcs_permissions]
 
 .. _howto/operator:CloudSQLNoOperationInProgressSensor:
 

providers/google/docs/operators/cloud/gcs.rst

@@ -87,6 +87,41 @@ More information
 See Google Cloud Storage Documentation to `create a new ACL entry for a bucket
 <https://cloud.google.com/storage/docs/json_api/v1/bucketAccessControls/insert>`_.
 
+.. _howto/operator:GCSBucketAddIamBindingOperator:
+
+GCSBucketAddIamBindingOperator
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Adds a member to an IAM role binding on the specified bucket. Unlike ACL operators,
+this operator works with uniform bucket-level access.
+
+For parameter definition, take a look at
+:class:`~airflow.providers.google.cloud.operators.gcs.GCSBucketAddIamBindingOperator`
+
+Using the operator
+""""""""""""""""""
+
+.. exampleinclude:: /../../google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py
+    :language: python
+    :dedent: 4
+    :start-after: [START howto_operator_cloudsql_export_gcs_permissions]
+    :end-before: [END howto_operator_cloudsql_export_gcs_permissions]
+
+Templating
+""""""""""
+
+.. literalinclude:: /../../google/src/airflow/providers/google/cloud/operators/gcs.py
+    :language: python
+    :dedent: 4
+    :start-after: [START gcs_bucket_add_iam_binding_template_fields]
+    :end-before: [END gcs_bucket_add_iam_binding_template_fields]
+
+More information
+""""""""""""""""
+
+See Google Cloud Storage documentation to `set and manage IAM policies on buckets
+<https://cloud.google.com/storage/docs/access-control/using-iam-permissions>`_.
+
 .. _howto/operator:GCSObjectCreateAclEntryOperator:
 
 GCSObjectCreateAclEntryOperator

providers/google/src/airflow/providers/google/cloud/hooks/gcs.py

@@ -1201,6 +1201,36 @@ def insert_bucket_acl(
 
         self.log.info("A new ACL entry created in bucket: %s", bucket_name)
 
+    def add_bucket_iam_binding(
+        self,
+        bucket_name: str,
+        role: str,
+        member: str,
+        user_project: str | None = None,
+    ) -> None:
+        """
+        Add a member to an IAM role binding on a bucket.
+
+        :param bucket_name: Name of a bucket.
+        :param role: The IAM role to grant.
+        :param member: The IAM member to grant the role to.
+        :param user_project: (Optional) The project to be billed for this request.
+            Required for Requester Pays buckets.
+        """
+        self.log.info("Adding %s to IAM role %s on bucket %s", member, role, bucket_name)
+        client = self.get_conn()
+        bucket = client.bucket(bucket_name=bucket_name, user_project=user_project)
+        policy = bucket.get_iam_policy(requested_policy_version=3)
+        for binding in policy.bindings:
+            if binding["role"] == role and binding.get("condition") is None:
+                binding["members"].add(member)
+                break
+        else:
+            policy.bindings.append({"role": role, "members": {member}})
+        bucket.set_iam_policy(policy)
+
+        self.log.info("Added %s to IAM role %s on bucket %s", member, role, bucket_name)
+
     def insert_object_acl(
         self,
         bucket_name: str,

providers/google/src/airflow/providers/google/cloud/operators/gcs.py

@@ -36,7 +36,7 @@
 if TYPE_CHECKING:
     from airflow.providers.common.compat.sdk import Context
 
-from google.api_core.exceptions import Conflict
+from google.api_core.exceptions import Conflict, GoogleAPIError
 from google.cloud.exceptions import GoogleCloudError
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -456,6 +456,91 @@ def execute(self, context: Context) -> None:
         )
 
 
+class GCSBucketAddIamBindingOperator(GoogleCloudBaseOperator):
+    """
+    Adds a member to an IAM role binding on the specified bucket.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:GCSBucketAddIamBindingOperator`
+
+    :param bucket: Name of a bucket.
+    :param role: The IAM role to grant, for example ``roles/storage.objectViewer``.
+    :param member: The IAM member to grant the role to, for example
+        ``serviceAccount:example@example-project.iam.gserviceaccount.com``.
+    :param user_project: (Optional) The project to be billed for this request.
+        Required for Requester Pays buckets.
+    :param gcp_conn_id: (Optional) The connection ID used to connect to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    # [START gcs_bucket_add_iam_binding_template_fields]
+    template_fields: Sequence[str] = (
+        "bucket",
+        "role",
+        "member",
+        "user_project",
+        "gcp_conn_id",
+        "impersonation_chain",
+    )
+    # [END gcs_bucket_add_iam_binding_template_fields]
+    operator_extra_links = (StorageLink(),)
+
+    def __init__(
+        self,
+        *,
+        bucket: str,
+        role: str,
+        member: str,
+        user_project: str | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.bucket = bucket
+        self.role = role
+        self.member = member
+        self.user_project = user_project
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def execute(self, context: Context) -> None:
+        hook = GCSHook(
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+        StorageLink.persist(
+            context=context,
+            uri=self.bucket,
+            project_id=hook.project_id,
+        )
+        try:
+            hook.add_bucket_iam_binding(
+                bucket_name=self.bucket,
+                role=self.role,
+                member=self.member,
+                user_project=self.user_project,
+            )
+        except GoogleAPIError as e:
+            self.log.exception(
+                "Failed to add member %s to IAM role %s on bucket %s. Google Cloud API error (%s): %s",
+                self.member,
+                self.role,
+                self.bucket,
+                type(e).__name__,
+                e,
+            )
+            raise
+
+
 class GCSObjectCreateAclEntryOperator(GoogleCloudBaseOperator):
     """
     Creates a new ACL entry on the specified object.

providers/google/tests/system/google/cloud/cloud_memorystore/example_cloud_memorystore_redis.py

@@ -41,7 +41,7 @@
     CloudMemorystoreUpdateInstanceOperator,
 )
 from airflow.providers.google.cloud.operators.gcs import (
-    GCSBucketCreateAclEntryOperator,
+    GCSBucketAddIamBindingOperator,
     GCSCreateBucketOperator,
     GCSDeleteBucketOperator,
 )
@@ -84,7 +84,9 @@
     tags=["example"],
 ) as dag:
     create_bucket = GCSCreateBucketOperator(
-        task_id="create_bucket", bucket_name=BUCKET_NAME, resource={"predefined_acl": "public_read_write"}
+        task_id="create_bucket",
+        bucket_name=BUCKET_NAME,
+        resource={"iamConfiguration": {"uniformBucketLevelAccess": {"enabled": True}}},
     )
 
     # [START howto_operator_create_instance]
@@ -163,15 +165,23 @@
     )
     # [END howto_operator_update_instance]
 
-    # [START howto_operator_set_acl_permission]
-    set_acl_permission = GCSBucketCreateAclEntryOperator(
-        task_id="gcs-set-acl-permission",
+    # [START howto_operator_set_iam_permissions]
+    redis_service_account = (
+        "{{ task_instance.xcom_pull(task_ids='get-instance')['persistence_iam_identity'] }}"
+    )
+    set_bucket_reader_permission = GCSBucketAddIamBindingOperator(
+        task_id="gcs-set-bucket-reader-permission",
+        bucket=BUCKET_NAME,
+        role="roles/storage.legacyBucketReader",
+        member=redis_service_account,
+    )
+    set_object_admin_permission = GCSBucketAddIamBindingOperator(
+        task_id="gcs-set-object-admin-permission",
         bucket=BUCKET_NAME,
-        entity="user-{{ task_instance.xcom_pull('get-instance')['persistence_iam_identity']"
-        ".split(':', 2)[1] }}",
-        role="OWNER",
+        role="roles/storage.objectAdmin",
+        member=redis_service_account,
     )
-    # [END howto_operator_set_acl_permission]
+    # [END howto_operator_set_iam_permissions]
 
     # [START howto_operator_export_instance]
     export_instance = CloudMemorystoreExportInstanceOperator(
@@ -253,7 +263,8 @@
         >> create_instance_result
         >> get_instance
         >> get_instance_result
-        >> set_acl_permission
+        >> set_bucket_reader_permission
+        >> set_object_admin_permission
         >> export_instance
         >> update_instance
         >> list_instances

providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py

@@ -25,7 +25,6 @@
 
 import os
 from datetime import datetime
-from urllib.parse import urlsplit
 
 from airflow.models.dag import DAG
 from airflow.models.xcom_arg import XComArg
@@ -41,10 +40,9 @@
     CloudSQLPatchInstanceDatabaseOperator,
 )
 from airflow.providers.google.cloud.operators.gcs import (
-    GCSBucketCreateAclEntryOperator,
+    GCSBucketAddIamBindingOperator,
     GCSCreateBucketOperator,
     GCSDeleteBucketOperator,
-    GCSObjectCreateAclEntryOperator,
 )
 from airflow.providers.google.cloud.sensors.cloud_sql import CloudSQLNoOperationInProgressSensor
 
@@ -146,7 +144,9 @@
     tags=["example", "cloud_sql"],
 ) as dag:
     create_bucket = GCSCreateBucketOperator(
-        task_id="create_bucket", bucket_name=BUCKET_NAME, resource={"predefined_acl": "public_read_write"}
+        task_id="create_bucket",
+        bucket_name=BUCKET_NAME,
+        resource={"iamConfiguration": {"uniformBucketLevelAccess": {"enabled": True}}},
     )
 
     # ############################################## #
@@ -187,18 +187,16 @@
     # ############################################## #
     # ### EXPORTING & IMPORTING SQL ################ #
     # ############################################## #
-    file_url_split = urlsplit(FILE_URI)
-
     # For export & import to work we need to add the Cloud SQL instance's Service Account
-    # write access to the destination GCS bucket.
+    # object admin access to the GCS bucket.
     service_account_email = XComArg(sql_instance_create_task, key="service_account_email")
 
     # [START howto_operator_cloudsql_export_gcs_permissions]
-    sql_gcp_add_bucket_permission_task = GCSBucketCreateAclEntryOperator(
-        entity=f"user-{service_account_email}",
-        role="WRITER",
-        bucket=file_url_split[1],  # netloc (bucket)
+    sql_gcp_add_bucket_permission_task = GCSBucketAddIamBindingOperator(
         task_id="sql_gcp_add_bucket_permission_task",
+        bucket=BUCKET_NAME,
+        role="roles/storage.objectAdmin",
+        member=f"serviceAccount:{service_account_email}",
     )
     # [END howto_operator_cloudsql_export_gcs_permissions]
 
@@ -217,18 +215,6 @@
     )
     # [END howto_operator_cloudsql_export_async]
 
-    # For import to work we need to add the Cloud SQL instance's Service Account
-    # read access to the target GCS object.
-    # [START howto_operator_cloudsql_import_gcs_permissions]
-    sql_gcp_add_object_permission_task = GCSObjectCreateAclEntryOperator(
-        entity=f"user-{service_account_email}",
-        role="READER",
-        bucket=file_url_split[1],  # netloc (bucket)
-        object_name=file_url_split[2][1:],  # path (strip first '/')
-        task_id="sql_gcp_add_object_permission_task",
-    )
-    # [END howto_operator_cloudsql_import_gcs_permissions]
-
     # Cloud SQL serializes admin operations per instance, so wait until the export above has
     # finished (no operation in progress) before submitting the import to avoid a 409.
     # [START howto_sensor_cloudsql_no_operation_in_progress]
@@ -297,7 +283,6 @@
         >> sql_gcp_add_bucket_permission_task
         >> sql_export_task
         >> sql_export_def_task
-        >> sql_gcp_add_object_permission_task
         >> sql_wait_no_operation_task
         >> sql_import_task
         >> sql_instance_clone