From f5f20a7b5c0d4ea0fa1e0a407d7e532b29939f62 Mon Sep 17 00:00:00 2001 From: DEEPIKA-217M <24uec125deepika@kgkite.ac.in> Date: Fri, 10 Jul 2026 09:30:28 +0530 Subject: [PATCH 1/2] Add Google Workspace Domain-Wide Delegation subject support --- .../google/cloud/transfers/gdrive_to_gcs.py | 3 +++ .../google/cloud/transfers/gdrive_to_local.py | 3 +++ .../cloud/utils/credentials_provider.py | 15 ++++++++++----- .../google/common/hooks/base_google.py | 19 ++++++++++++------- .../google/common/hooks/discovery_api.py | 16 ++++++++++------ .../providers/google/suite/hooks/drive.py | 4 ++++ .../providers/google/suite/hooks/sheets.py | 16 ++++++++++------ .../providers/google/suite/sensors/drive.py | 3 +++ .../google/suite/transfers/gcs_to_gdrive.py | 3 +++ .../google/suite/transfers/local_to_drive.py | 3 +++ 10 files changed, 61 insertions(+), 24 deletions(-) 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..100d40f62fc50 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 7ffef424bf83d..30dea6ac756dd 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 @@ -299,19 +299,22 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: } def __init__( - self, - gcp_conn_id: str = "google_cloud_default", - impersonation_chain: str | Sequence[str] | None = None, - *, - quota_project_id: str | None = None, - **kwargs, - ) -> None: + self, + 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: """ Initialize the Google Cloud Base Hook. :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..3cc7b17d136a7 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,20 +42,24 @@ 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 def __init__( - self, - api_service_name: str, - api_version: str, - gcp_conn_id: str = "google_cloud_default", - impersonation_chain: str | Sequence[str] | None = None, - ) -> None: + self, + api_service_name: str, + 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..783f760481a10 100644 --- a/providers/google/src/airflow/providers/google/suite/hooks/sheets.py +++ b/providers/google/src/airflow/providers/google/suite/hooks/sheets.py @@ -46,18 +46,22 @@ 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__( - self, - gcp_conn_id: str = "google_cloud_default", - api_version: str = "v4", - impersonation_chain: str | Sequence[str] | None = None, - api_endpoint: str | None = None, - ) -> None: + self, + 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..97335f8f4fd9d 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 = [] From fe5344aa7307167f9e7150bf51c8f5c29bc0a21d Mon Sep 17 00:00:00 2001 From: DEEPIKA-217M <24uec125deepika@kgkite.ac.in> Date: Tue, 14 Jul 2026 09:14:08 +0530 Subject: [PATCH 2/2] Add Google Workspace Domain-Wide Delegation subject support --- .../google/cloud/utils/credentials_provider.py | 4 ++-- .../providers/google/common/hooks/base_google.py | 16 ++++++++-------- .../google/common/hooks/discovery_api.py | 14 +++++++------- .../providers/google/suite/hooks/sheets.py | 16 ++++++++-------- .../google/suite/transfers/local_to_drive.py | 2 +- 5 files changed, 26 insertions(+), 26 deletions(-) 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 100d40f62fc50..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 @@ -277,7 +277,7 @@ 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() - + delegated_subject = self.subject or self.delegate_to if delegated_subject: @@ -287,7 +287,7 @@ def get_credentials_and_project(self) -> tuple[Credentials, str]: raise AirflowException( "The current authentication method does not support " "Google Workspace Domain-Wide Delegation." - "Please use service account credentials." + "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 30dea6ac756dd..c12eda4dd40fa 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 @@ -299,14 +299,14 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]: } def __init__( - self, - 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: + self, + 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: """ Initialize the Google Cloud Base Hook. 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 3cc7b17d136a7..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 @@ -49,13 +49,13 @@ class GoogleDiscoveryApiHook(GoogleBaseHook): _conn: Resource | None = None def __init__( - self, - api_service_name: str, - api_version: str, - gcp_conn_id: str = "google_cloud_default", - impersonation_chain: str | Sequence[str] | None = None, - subject: str | None = None, -) -> None: + self, + api_service_name: str, + 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, 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 783f760481a10..94d35eb2a00a2 100644 --- a/providers/google/src/airflow/providers/google/suite/hooks/sheets.py +++ b/providers/google/src/airflow/providers/google/suite/hooks/sheets.py @@ -51,17 +51,17 @@ class GSheetsHook(GoogleBaseHook): """ def __init__( - self, - 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: + self, + 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, + subject=subject, ) self.gcp_conn_id = gcp_conn_id self.api_version = api_version 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 97335f8f4fd9d..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 @@ -108,7 +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, + subject=self.subject, ) remote_file_ids = []