Skip to content

Implement token refresh on BigQuery dts run operator to handle the heavy migration jobs over 1 hour #37538

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
okayhooni wants to merge 1 commit into from
Closed

Implement token refresh on BigQuery dts run operator to handle the heavy migration jobs over 1 hour #37538

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 23 additions & 1 deletion airflow/providers/google/cloud/operators/bigquery_dts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand Down Expand Up @@ -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] = (
Expand All @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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",
)
Expand All @@ -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"
)
Comment on lines +374 to +382

@eladkal eladkal Feb 19, 2024

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Other providers also requires refresh logic but seems like they have simpler implementation. For example:

airflow/airflow/providers/databricks/hooks/databricks_base.py

Lines 441 to 453 in 0b680c9

@staticmethod
def _is_oauth_token_valid(token: dict, time_key="expires_on") -> bool:
"""
Check if an OAuth token is valid and hasn't expired yet.
:param sp_token: dict with properties of OAuth token
:param time_key: name of the key that holds the time of expiration
:return: true if token is valid, false otherwise
"""
if "access_token" not in token or token.get("token_type", "") != "Bearer" or time_key not in token:
raise AirflowException(f"Can't get necessary data from OAuth token: {token}")
return int(token[time_key]) > (int(time.time()) + TOKEN_REFRESH_LEAD_TIME)

Can you explain why we need complex logic here?

Also please extract the check validation (time handling) to it's own private function - it should be covered by unit tests

@okayhooni okayhooni Feb 19, 2024 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you for quick and detailed review..!

Actually, I can't find the token validation method on base_google.py module, as like databricks_base.py module you mentioned.. If it has similar method based on Google sdk, it will be more simple and elegant way, as you said. But, it tasks more GCP API calls to check token validity on every loop (+ Is it possible to check token validity or refresh token with already expired credential token..?)

We experienced some error cases, even though the tasks take just over 30 minutes, not 1 hour, with same authorization error. (I don't know why..) so I decided to use the half of default token lifespan to refresh token.


transfer_run: TransferRun = self.hook.get_transfer_run(
run_id=run_id,
transfer_config_id=transfer_config_id,
Expand All @@ -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:
Expand Down
22 changes: 22 additions & 0 deletions airflow/providers/google/cloud/triggers/bigquery_dts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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__(
Expand All @@ -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
Expand All @@ -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."""
Expand All @@ -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,
Expand Down Expand Up @@ -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(
{
Expand Down
11 changes: 11 additions & 0 deletions airflow/providers/google/common/hooks/base_google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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)()