diff --git a/airflow/providers/google/cloud/operators/bigquery_dts.py b/airflow/providers/google/cloud/operators/bigquery_dts.py index 50948a0614a2c..d9e20ab65b639 100644 --- a/airflow/providers/google/cloud/operators/bigquery_dts.py +++ b/airflow/providers/google/cloud/operators/bigquery_dts.py @@ -19,6 +19,7 @@ from __future__ import annotations import time +from datetime import timedelta from functools import cached_property from typing import TYPE_CHECKING, Sequence @@ -29,6 +30,7 @@ TransferRun, TransferState, ) +from google.oauth2.service_account import _DEFAULT_TOKEN_LIFETIME_SECS from airflow.configuration import conf from airflow.exceptions import AirflowException @@ -256,6 +258,7 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat 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). :param deferrable: Run operator in the deferrable mode. + :param token_refresh_interval_seconds: GCP STS Token refresh interval in seconds. """ template_fields: Sequence[str] = ( @@ -282,6 +285,7 @@ def __init__( gcp_conn_id="google_cloud_default", impersonation_chain: str | Sequence[str] | None = None, deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + token_refresh_interval_seconds: int = _DEFAULT_TOKEN_LIFETIME_SECS // 2, **kwargs, ) -> None: super().__init__(**kwargs) @@ -291,11 +295,16 @@ def __init__( self.requested_time_range = requested_time_range self.requested_run_time = requested_run_time self.retry = retry - self.timeout = timeout self.metadata = metadata self.gcp_conn_id = gcp_conn_id self.impersonation_chain = impersonation_chain self.deferrable = deferrable + self.token_refresh_interval_seconds = token_refresh_interval_seconds + self.timeout = ( + (self.execution_timeout + timedelta(minutes=10)).total_seconds() + if self.execution_timeout and not timeout + else timeout + ) @cached_property def hook(self) -> BiqQueryDataTransferServiceHook: @@ -351,6 +360,7 @@ def execute(self, context: Context): gcp_conn_id=self.gcp_conn_id, location=self.location, impersonation_chain=self.impersonation_chain, + token_refresh_interval_seconds=self.token_refresh_interval_seconds, ), method_name="execute_completed", ) @@ -359,7 +369,18 @@ def _wait_for_transfer_to_be_done(self, run_id: str, transfer_config_id: str, in if interval <= 0: raise ValueError("Interval must be > 0") + idx = 0 while True: + current_tick_div, current_tick_mod = divmod(idx * interval, self.token_refresh_interval_seconds) + next_tick_div, next_tick_mod = divmod((idx + 1) * interval, self.token_refresh_interval_seconds) + if (current_tick_div < next_tick_div and 0 < next_tick_mod) or ( + current_tick_mod == 0 and 1 <= current_tick_div + ): + _ = self.hook.refresh_credentials() + self.log.info( + f"Credentials were refreshed on tick: idx={idx}, idx*interval={idx * interval} sec" + ) + transfer_run: TransferRun = self.hook.get_transfer_run( run_id=run_id, transfer_config_id=transfer_config_id, @@ -380,6 +401,7 @@ def _wait_for_transfer_to_be_done(self, run_id: str, transfer_config_id: str, in self.log.info("Transfer run is still working, waiting for %s seconds...", interval) self.log.info("Transfer run status: %s", state) time.sleep(interval) + idx += 1 @staticmethod def _job_is_done(state: TransferState) -> bool: diff --git a/airflow/providers/google/cloud/triggers/bigquery_dts.py b/airflow/providers/google/cloud/triggers/bigquery_dts.py index 4d4d7fbdfa892..f39cb41354092 100644 --- a/airflow/providers/google/cloud/triggers/bigquery_dts.py +++ b/airflow/providers/google/cloud/triggers/bigquery_dts.py @@ -21,6 +21,7 @@ from typing import Any, AsyncIterator, Sequence from google.cloud.bigquery_datatransfer_v1 import TransferRun, TransferState +from google.oauth2.service_account import _DEFAULT_TOKEN_LIFETIME_SECS from airflow.providers.google.cloud.hooks.bigquery_dts import AsyncBiqQueryDataTransferServiceHook from airflow.triggers.base import BaseTrigger, TriggerEvent @@ -44,6 +45,7 @@ class BigQueryDataTransferRunTrigger(BaseTrigger): 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). + :param token_refresh_interval_seconds: GCP STS Token refresh interval in seconds. """ def __init__( @@ -55,6 +57,7 @@ def __init__( gcp_conn_id: str = "google_cloud_default", location: str | None = None, impersonation_chain: str | Sequence[str] | None = None, + token_refresh_interval_seconds: int = _DEFAULT_TOKEN_LIFETIME_SECS // 2, ): super().__init__() self.project_id = project_id @@ -64,6 +67,7 @@ def __init__( self.gcp_conn_id = gcp_conn_id self.location = location self.impersonation_chain = impersonation_chain + self.token_refresh_interval_seconds = token_refresh_interval_seconds def serialize(self) -> tuple[str, dict[str, Any]]: """Serialize class arguments and classpath.""" @@ -77,14 +81,30 @@ def serialize(self) -> tuple[str, dict[str, Any]]: "gcp_conn_id": self.gcp_conn_id, "location": self.location, "impersonation_chain": self.impersonation_chain, + "token_refresh_interval_seconds": self.token_refresh_interval_seconds, }, ) async def run(self) -> AsyncIterator[TriggerEvent]: """If the Transfer Run is in a terminal state, then yield TriggerEvent object.""" hook = self._get_async_hook() + idx = 0 try: while True: + current_tick_div, current_tick_mod = divmod( + idx * self.poll_interval, self.token_refresh_interval_seconds + ) + next_tick_div, next_tick_mod = divmod( + (idx + 1) * self.poll_interval, self.token_refresh_interval_seconds + ) + if (current_tick_div < next_tick_div and 0 < next_tick_mod) or ( + current_tick_mod == 0 and 1 <= current_tick_div + ): + _ = await hook.refresh_credentials() + self.log.info( + f"Credentials were refreshed on tick: idx={idx}, idx*interval={idx * self.poll_interval} sec" + ) + transfer_run: TransferRun = await hook.get_transfer_run( project_id=self.project_id, config_id=self.config_id, @@ -129,6 +149,8 @@ async def run(self) -> AsyncIterator[TriggerEvent]: self.log.info("Job is still working...") self.log.info("Waiting for %s seconds", self.poll_interval) await asyncio.sleep(self.poll_interval) + + idx += 1 except Exception as e: yield TriggerEvent( { diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index bee2798267610..fffd6167bf654 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -307,6 +307,12 @@ def get_credentials(self) -> google.auth.credentials.Credentials: credentials, _ = self.get_credentials_and_project_id() return credentials + def refresh_credentials(self) -> google.auth.credentials.Credentials: + """Refresh the Credentials object for Google API.""" + credentials = self.get_credentials() + credentials.refresh(google.auth.transport.requests.Request()) + return credentials + def _get_access_token(self) -> str: """Return a valid access token from Google API Credentials.""" credentials = self.get_credentials() @@ -706,3 +712,8 @@ async def service_file_as_context(self) -> Any: """ sync_hook = await self.get_sync_hook() return await sync_to_async(sync_hook.provide_gcp_credential_file_as_context)() + + async def refresh_credentials(self) -> google.auth.credentials.Credentials: + """Refresh the Credentials object for Google API.""" + sync_hook = await self.get_sync_hook() + return await sync_to_async(sync_hook.refresh_credentials)()