diff --git a/providers/google/docs/operators/cloud/cloud_memorystore.rst b/providers/google/docs/operators/cloud/cloud_memorystore.rst index c03e9a876c150..c51dc94cf43d2 100644 --- a/providers/google/docs/operators/cloud/cloud_memorystore.rst +++ b/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 -`__ +`__ .. _howto/operator:CloudMemorystoreCreateInstanceOperator: diff --git a/providers/google/docs/operators/cloud/cloud_sql.rst b/providers/google/docs/operators/cloud/cloud_sql.rst index 544ab2071579c..e40f9e4254f67 100644 --- a/providers/google/docs/operators/cloud/cloud_sql.rst +++ b/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: diff --git a/providers/google/docs/operators/cloud/gcs.rst b/providers/google/docs/operators/cloud/gcs.rst index 2d7910ff31b80..4fc9f17293f62 100644 --- a/providers/google/docs/operators/cloud/gcs.rst +++ b/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 `_. +.. _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 +`_. + .. _howto/operator:GCSObjectCreateAclEntryOperator: GCSObjectCreateAclEntryOperator diff --git a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py b/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py index e9eb0666cb9bb..0693be10d96c0 100644 --- a/providers/google/src/airflow/providers/google/cloud/hooks/gcs.py +++ b/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, diff --git a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py b/providers/google/src/airflow/providers/google/cloud/operators/gcs.py index b55709c309eea..9c0d9758a792a 100644 --- a/providers/google/src/airflow/providers/google/cloud/operators/gcs.py +++ b/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. diff --git a/providers/google/tests/system/google/cloud/cloud_memorystore/example_cloud_memorystore_redis.py b/providers/google/tests/system/google/cloud/cloud_memorystore/example_cloud_memorystore_redis.py index 3de2140b74ecf..95b5c39ce1cbb 100644 --- a/providers/google/tests/system/google/cloud/cloud_memorystore/example_cloud_memorystore_redis.py +++ b/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 diff --git a/providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py b/providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py index eda31bd43c9a6..9c0db759521d4 100644 --- a/providers/google/tests/system/google/cloud/cloud_sql/example_cloud_sql.py +++ b/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 diff --git a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py index 59bf1da4f0baa..e1a9066740f27 100644 --- a/providers/google/tests/unit/google/cloud/hooks/test_gcs.py +++ b/providers/google/tests/unit/google/cloud/hooks/test_gcs.py @@ -896,6 +896,67 @@ def test_create_bucket_with_resource(self, mock_service, mock_bucket): project=test_project, location=test_location ) + @pytest.mark.parametrize( + ("bindings", "expected_bindings"), + [ + ( + [], + [{"role": "roles/storage.objectViewer", "members": {"serviceAccount:test@example.com"}}], + ), + ( + [{"role": "roles/storage.objectViewer", "members": {"user:existing@example.com"}}], + [ + { + "role": "roles/storage.objectViewer", + "members": { + "serviceAccount:test@example.com", + "user:existing@example.com", + }, + } + ], + ), + ( + [ + { + "role": "roles/storage.objectViewer", + "members": {"user:conditional@example.com"}, + "condition": {"title": "conditional-binding"}, + } + ], + [ + { + "role": "roles/storage.objectViewer", + "members": {"user:conditional@example.com"}, + "condition": {"title": "conditional-binding"}, + }, + { + "role": "roles/storage.objectViewer", + "members": {"serviceAccount:test@example.com"}, + }, + ], + ), + ], + ) + @mock.patch(GCS_STRING.format("GCSHook.get_conn")) + def test_add_bucket_iam_binding(self, mock_service, bindings, expected_bindings): + bucket = mock_service.return_value.bucket.return_value + policy = mock.MagicMock(bindings=copy.deepcopy(bindings)) + bucket.get_iam_policy.return_value = policy + + self.gcs_hook.add_bucket_iam_binding( + bucket_name="test-bucket", + role="roles/storage.objectViewer", + member="serviceAccount:test@example.com", + user_project="test-user-project", + ) + + mock_service.return_value.bucket.assert_called_once_with( + bucket_name="test-bucket", user_project="test-user-project" + ) + bucket.get_iam_policy.assert_called_once_with(requested_policy_version=3) + assert policy.bindings == expected_bindings + bucket.set_iam_policy.assert_called_once_with(policy) + @mock.patch("google.cloud.storage.Bucket.blob") @mock.patch(GCS_STRING.format("GCSHook.get_conn")) def test_compose(self, mock_service, mock_blob): diff --git a/providers/google/tests/unit/google/cloud/operators/test_gcs.py b/providers/google/tests/unit/google/cloud/operators/test_gcs.py index 664308cd29188..44ceea9aa8cff 100644 --- a/providers/google/tests/unit/google/cloud/operators/test_gcs.py +++ b/providers/google/tests/unit/google/cloud/operators/test_gcs.py @@ -23,6 +23,7 @@ from unittest import mock import pytest +from google.api_core.exceptions import RetryError from google.cloud.exceptions import GoogleCloudError from airflow.exceptions import AirflowProviderDeprecationWarning @@ -33,6 +34,7 @@ PreviousIdentifier, ) from airflow.providers.google.cloud.operators.gcs import ( + GCSBucketAddIamBindingOperator, GCSBucketCreateAclEntryOperator, GCSCreateBucketOperator, GCSDeleteBucketOperator, @@ -120,6 +122,61 @@ def test_object_create_acl(self, mock_hook): ) +class TestGoogleCloudStorageIam: + @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook") + def test_add_bucket_iam_binding(self, mock_hook): + operator = GCSBucketAddIamBindingOperator( + bucket="test-bucket", + role="roles/storage.objectViewer", + member="serviceAccount:test@example.com", + user_project="test-user-project", + task_id="id", + ) + + operator.execute(context=mock.MagicMock()) + + mock_hook.return_value.add_bucket_iam_binding.assert_called_once_with( + bucket_name="test-bucket", + role="roles/storage.objectViewer", + member="serviceAccount:test@example.com", + user_project="test-user-project", + ) + + @pytest.mark.parametrize( + "error", + [ + GoogleCloudError("test API error"), + RetryError("retries exhausted", GoogleCloudError("last API error")), + ], + ids=["api-error", "retry-error"], + ) + @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook") + def test_add_bucket_iam_binding_reraises_api_error(self, mock_hook, error): + mock_hook.return_value.add_bucket_iam_binding.side_effect = error + operator = GCSBucketAddIamBindingOperator( + bucket="test-bucket", + role="roles/storage.objectViewer", + member="serviceAccount:test@example.com", + task_id="id", + ) + + with ( + mock.patch.object(operator.log, "exception") as mock_log_exception, + pytest.raises(type(error)) as raised, + ): + operator.execute(context=mock.MagicMock()) + + assert raised.value is error + mock_log_exception.assert_called_once_with( + "Failed to add member %s to IAM role %s on bucket %s. Google Cloud API error (%s): %s", + "serviceAccount:test@example.com", + "roles/storage.objectViewer", + "test-bucket", + type(error).__name__, + error, + ) + + class TestGCSDeleteObjectsOperator: @mock.patch("airflow.providers.google.cloud.operators.gcs.GCSHook") def test_delete_objects(self, mock_hook):