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

[okayhooni]

related issue:

• Getting Unauthorized error messages with GKEStartPodOperator if pod execution is over 1 hour #31648

releated PR:

• Support providing creds lifetime for GCP connection token #31651
• Refresh GKE OAuth2 tokens #32673

Currently, when BigQuery Data Transfer job triggered by Airflow operator, exceeds one hour, it fails due to the expiration of the credential (default lifespan = 1 hour).

Access Denied: Table projectid:table_name: Permission bigquery.tables.get denied on table projectid:datasetid.table_name (or it may not exist).

So, I add the token refresh logic on DTS run operator. (additionally, add fallback logic to timeout parameter based on execution_timeout of Operator)

[eladkal]

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]

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.

[okayhooni]

I got the same authorization error even after applying this refresh logic.. So, I close it.

[eladkal]

cc @VladaZakharova maybe something worth investigating