Add subject support for Google Workspace Domain-Wide Delegation

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,

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

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:

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,

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

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
 

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,
+        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

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)

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:

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,13 +100,15 @@ 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
 
     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 = []