diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py index 6cd1808c79ca1..8be8ab1d13c65 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_gcs.py @@ -72,6 +72,7 @@ def __init__( drive_id: str | None = None, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -82,11 +83,13 @@ def __init__( self.file_name = file_name self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain + self.subject = subject def execute(self, context: Context) -> list[str]: gdrive_hook = GoogleDriveHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + subject=self.subject, ) gcs_hook = GCSHook( gcp_conn_id=self.gcp_conn_id, diff --git a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py index e7fe7d731eb77..623756393b0a2 100644 --- a/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py +++ b/providers/google/src/airflow/providers/google/cloud/transfers/gdrive_to_local.py @@ -66,6 +66,7 @@ def __init__( drive_id: str | None = None, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -75,12 +76,14 @@ def __init__( self.file_name = file_name self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain + self.subject = subject def execute(self, context: Context): self.log.info("Executing download: %s into %s", self.file_name, self.output_file) gdrive_hook = GoogleDriveHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + subject=self.subject, ) file_metadata = gdrive_hook.get_file_id( folder_id=self.folder_id, file_name=self.file_name, drive_id=self.drive_id diff --git a/providers/google/src/airflow/providers/google/cloud/utils/credentials_provider.py b/providers/google/src/airflow/providers/google/cloud/utils/credentials_provider.py index b363d233e6ec3..c3b1a4a0504cb 100644 --- a/providers/google/src/airflow/providers/google/cloud/utils/credentials_provider.py +++ b/providers/google/src/airflow/providers/google/cloud/utils/credentials_provider.py @@ -209,6 +209,7 @@ def __init__( key_secret_project_id: str | None = None, scopes: Collection[str] | None = None, delegate_to: str | None = None, + subject: str | None = None, disable_logging: bool = False, target_principal: str | None = None, delegates: Sequence[str] | None = None, @@ -240,6 +241,7 @@ def __init__( self.key_secret_project_id = key_secret_project_id self.scopes = scopes self.delegate_to = delegate_to + self.subject = subject self.disable_logging = disable_logging self.target_principal = target_principal self.delegates = delegates @@ -275,14 +277,17 @@ def get_credentials_and_project(self) -> tuple[Credentials, str]: credentials, project_id = self._get_credentials_using_credential_config_file() else: credentials, project_id = self._get_credentials_using_adc() - if self.delegate_to: + + delegated_subject = self.subject or self.delegate_to + + if delegated_subject: if hasattr(credentials, "with_subject"): - credentials = credentials.with_subject(self.delegate_to) + credentials = credentials.with_subject(delegated_subject) else: raise AirflowException( - "The `delegate_to` parameter cannot be used here as the current " - "authentication method does not support account impersonate. " - "Please use service-account for authorization." + "The current authentication method does not support " + "Google Workspace Domain-Wide Delegation." + "Please use service account credentials." ) if self.target_principal: diff --git a/providers/google/src/airflow/providers/google/common/hooks/base_google.py b/providers/google/src/airflow/providers/google/common/hooks/base_google.py index 649901971694b..a7ef60838dfe4 100644 --- a/providers/google/src/airflow/providers/google/common/hooks/base_google.py +++ b/providers/google/src/airflow/providers/google/common/hooks/base_google.py @@ -303,6 +303,7 @@ def __init__( gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, *, + subject: str | None = None, quota_project_id: str | None = None, **kwargs, ) -> None: @@ -312,6 +313,8 @@ def __init__( :param gcp_conn_id: The connection ID to use when fetching connection info. :param impersonation_chain: Optional service account to impersonate using short-term credentials. + :param subject: Optional Google Workspace user to impersonate using + Domain-Wide Delegation. :param quota_project_id: Optional Project ID to use for quota/billing purposes. If None, no separate quota project is configured and the default behavior of the credentials is used. @@ -320,6 +323,7 @@ def __init__( super().__init__(**kwargs) self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain + self.subject = subject if quota_project_id is not None: self._validate_quota_project(quota_project_id) self.quota_project_id = quota_project_id @@ -376,6 +380,7 @@ def get_credentials_and_project_id(self) -> tuple[Credentials, str | None]: key_secret_name=key_secret_name, key_secret_project_id=key_secret_project_id, scopes=self.scopes, + subject=self.subject, target_principal=target_principal, delegates=delegates, is_anonymous=is_anonymous, diff --git a/providers/google/src/airflow/providers/google/common/hooks/discovery_api.py b/providers/google/src/airflow/providers/google/common/hooks/discovery_api.py index 19c0fd53411bf..1abb4d6a325b0 100644 --- a/providers/google/src/airflow/providers/google/common/hooks/discovery_api.py +++ b/providers/google/src/airflow/providers/google/common/hooks/discovery_api.py @@ -42,6 +42,8 @@ class GoogleDiscoveryApiHook(GoogleBaseHook): 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. + :param subject: Optional Google Workspace user to impersonate using + Domain-Wide Delegation. """ _conn: Resource | None = None @@ -52,10 +54,12 @@ def __init__( api_version: str, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, ) -> None: super().__init__( gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, + subject=subject, ) self.api_service_name = api_service_name self.api_version = api_version diff --git a/providers/google/src/airflow/providers/google/suite/hooks/drive.py b/providers/google/src/airflow/providers/google/suite/hooks/drive.py index 7d1257cf4edc7..064130506e13a 100644 --- a/providers/google/src/airflow/providers/google/suite/hooks/drive.py +++ b/providers/google/src/airflow/providers/google/suite/hooks/drive.py @@ -43,6 +43,8 @@ class GoogleDriveHook(GoogleBaseHook): 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. + :param subject: Optional Google Workspace user to impersonate using + Domain-Wide Delegation. """ _conn: Resource | None = None @@ -52,10 +54,12 @@ def __init__( api_version: str = "v3", gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, ) -> None: super().__init__( gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, + subject=subject, ) self.api_version = api_version diff --git a/providers/google/src/airflow/providers/google/suite/hooks/sheets.py b/providers/google/src/airflow/providers/google/suite/hooks/sheets.py index 7c7ec98f0bd39..94d35eb2a00a2 100644 --- a/providers/google/src/airflow/providers/google/suite/hooks/sheets.py +++ b/providers/google/src/airflow/providers/google/suite/hooks/sheets.py @@ -46,6 +46,8 @@ class GSheetsHook(GoogleBaseHook): account from the list granting this role to the originating account. :param api_endpoint: Optional. Custom API endpoint, i.e: regional or private endpoint. This can be used to target private VPC or restricted access endpoints. + :param subject: Optional Google Workspace user to impersonate using + Domain-Wide Delegation. """ def __init__( @@ -53,11 +55,13 @@ def __init__( gcp_conn_id: str = "google_cloud_default", api_version: str = "v4", impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, api_endpoint: str | None = None, ) -> None: super().__init__( gcp_conn_id=gcp_conn_id, impersonation_chain=impersonation_chain, + subject=subject, ) self.gcp_conn_id = gcp_conn_id self.api_version = api_version diff --git a/providers/google/src/airflow/providers/google/suite/sensors/drive.py b/providers/google/src/airflow/providers/google/suite/sensors/drive.py index 14b4b1e4db06d..12ff8e8dde158 100644 --- a/providers/google/src/airflow/providers/google/suite/sensors/drive.py +++ b/providers/google/src/airflow/providers/google/suite/sensors/drive.py @@ -64,6 +64,7 @@ def __init__( drive_id: str | None = None, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -72,11 +73,13 @@ def __init__( self.drive_id = drive_id self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain + self.subject = subject def poke(self, context: Context) -> bool: self.log.info("Sensor is checking for the file %s in the folder %s", self.file_name, self.folder_id) hook = GoogleDriveHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + subject=self.subject, ) return hook.exists(folder_id=self.folder_id, file_name=self.file_name, drive_id=self.drive_id) diff --git a/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py b/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py index e184f00a1bc3f..12045ff7b0d37 100644 --- a/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py +++ b/providers/google/src/airflow/providers/google/suite/transfers/gcs_to_gdrive.py @@ -101,6 +101,7 @@ def __init__( move_object: bool = False, gcp_conn_id: str = "google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, **kwargs, ) -> None: super().__init__(**kwargs) @@ -112,6 +113,7 @@ def __init__( self.move_object = move_object self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain + self.subject = subject self.gcs_hook: GCSHook | None = None self.gdrive_hook: GoogleDriveHook | None = None @@ -123,6 +125,7 @@ def execute(self, context: Context): self.gdrive_hook = GoogleDriveHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + subject=self.subject, ) if WILDCARD in self.source_object: diff --git a/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py b/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py index 4c3b2c5fd0cc9..5a81e1f108a1d 100644 --- a/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py +++ b/providers/google/src/airflow/providers/google/suite/transfers/local_to_drive.py @@ -86,6 +86,7 @@ def __init__( chunk_size: int = 100 * 1024 * 1024, resumable: bool = False, impersonation_chain: str | Sequence[str] | None = None, + subject: str | None = None, folder_id: str = "root", show_full_target_path: bool = True, **kwargs, @@ -99,6 +100,7 @@ def __init__( self.chunk_size = chunk_size self.resumable = resumable self.impersonation_chain = impersonation_chain + self.subject = subject self.folder_id = folder_id self.show_full_target_path = show_full_target_path @@ -106,6 +108,7 @@ def execute(self, context: Context) -> list[str]: hook = GoogleDriveHook( gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain, + subject=self.subject, ) remote_file_ids = []