From a24d3f847684d426d7b97a866a63eaa553710a03 Mon Sep 17 00:00:00 2001 From: Maksim Moiseenkov Date: Fri, 2 Aug 2024 10:38:51 +0000 Subject: [PATCH] Upgrade package gcloud-aio-auth>=5.2.0 --- Dockerfile.ci | 2 +- .../google/common/hooks/base_google.py | 37 ++++++++++++++++--- airflow/providers/google/provider.yaml | 8 +--- .../apache-airflow-providers-google/index.rst | 2 +- generated/provider_dependencies.json | 2 +- scripts/docker/entrypoint_ci.sh | 2 +- 6 files changed, 36 insertions(+), 17 deletions(-) diff --git a/Dockerfile.ci b/Dockerfile.ci index e2042c668d6fb..1ef5d38389ffd 100644 --- a/Dockerfile.ci +++ b/Dockerfile.ci @@ -1035,7 +1035,7 @@ function check_boto_upgrade() { set -x # shellcheck disable=SC2086 ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade boto3 botocore \ - "oss2>=2.14.0" "gcloud-aio-auth>=4.0.0,<5.0.0" "requests!=2.32.*,<3.0.0,>=2.24.0" + "oss2>=2.14.0" "cryptography<43.0.0" "requests!=2.32.*,<3.0.0,>=2.24.0" set +x pip check } diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index a70922b5d3abb..a55bb49f766c2 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -19,6 +19,7 @@ from __future__ import annotations +import asyncio import datetime import functools import json @@ -36,7 +37,7 @@ import tenacity from asgiref.sync import sync_to_async from deprecated import deprecated -from gcloud.aio.auth.token import Token +from gcloud.aio.auth.token import Token, TokenResponse from google.api_core.exceptions import Forbidden, ResourceExhausted, TooManyRequests from google.auth import _cloud_sdk, compute_engine # type: ignore[attr-defined] from google.auth.environment_vars import CLOUD_SDK_CONFIG_DIR, CREDENTIALS @@ -745,17 +746,41 @@ async def from_hook( async def get_project(self) -> str | None: return self.project - async def acquire_access_token(self, timeout: int = 10) -> None: + async def refresh(self, *, timeout: int) -> TokenResponse: await sync_to_async(self.credentials.refresh)(google.auth.transport.requests.Request()) self.access_token = cast(str, self.credentials.token) self.access_token_duration = 3600 - # access_token_acquired_at is specific to gcloud-aio's Token. On subsequent calls of `get` it will be used - # with `datetime.datetime.utcnow()`. Therefore we have to use an offset-naive datetime. - # https://github.com/talkiq/gcloud-aio/blob/f1132b005ba35d8059229a9ca88b90f31f77456d/auth/gcloud/aio/auth/token.py#L204 - self.access_token_acquired_at = datetime.datetime.now(tz=datetime.timezone.utc).replace(tzinfo=None) + self.access_token_acquired_at = self._now() + return TokenResponse(value=self.access_token, expires_in=self.access_token_duration) + + async def acquire_access_token(self, timeout: int = 10) -> None: + await self.refresh(timeout=timeout) self.acquiring = None + async def ensure_token(self) -> None: + if self.acquiring and not self.acquiring.done(): + await self.acquiring + return + + if self.access_token: + delta = (self._now() - self.access_token_acquired_at).total_seconds() + if delta <= self.access_token_duration / 2: + return + + self.acquiring = asyncio.ensure_future( # pylint: disable=used-before-assignment + self.acquire_access_token() + ) + await self.acquiring + + @staticmethod + def _now(): + # access_token_acquired_at is specific to gcloud-aio's Token. + # On subsequent calls of `get` it will be used with `datetime.datetime.utcnow()`. + # Therefore we have to use an offset-naive datetime. + # https://github.com/talkiq/gcloud-aio/blob/f1132b005ba35d8059229a9ca88b90f31f77456d/auth/gcloud/aio/auth/token.py#L204 + return datetime.datetime.now(tz=datetime.timezone.utc).replace(tzinfo=None) + class GoogleBaseAsyncHook(BaseHook): """GoogleBaseAsyncHook inherits from BaseHook class, run on the trigger worker.""" diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml index 8b262ff7d4231..2d26522bb5062 100644 --- a/airflow/providers/google/provider.yaml +++ b/airflow/providers/google/provider.yaml @@ -98,13 +98,7 @@ dependencies: - apache-airflow-providers-common-sql>=1.7.2 - asgiref>=3.5.2 - dill>=0.2.3 - # When upgrading the major version of gcloud-aio-auth we want to make sure to - # 1. use at least version 5.2, which uses offset-aware datetime internally - # 2. override Token's new `refresh` method instead of `acquire_access_token`, which allows us to avoid - # dealing with internals like `access_token_acquired_at` - # 3. continue to `subclass gcloud.aio.auth.token.Token` instead of `BaseToken`, since instances of - # `_CredentialsToken` are instances of `Token` and used as such - - gcloud-aio-auth>=4.0.0,<5.0.0 + - gcloud-aio-auth>=5.2.0 - gcloud-aio-bigquery>=6.1.2 - gcloud-aio-storage>=9.0.0 - gcsfs>=2023.10.0 diff --git a/docs/apache-airflow-providers-google/index.rst b/docs/apache-airflow-providers-google/index.rst index 769eaf1011788..cdc6ba513381a 100644 --- a/docs/apache-airflow-providers-google/index.rst +++ b/docs/apache-airflow-providers-google/index.rst @@ -115,7 +115,7 @@ PIP package Version required ``apache-airflow-providers-common-sql`` ``>=1.7.2`` ``asgiref`` ``>=3.5.2`` ``dill`` ``>=0.2.3`` -``gcloud-aio-auth`` ``>=4.0.0,<5.0.0`` +``gcloud-aio-auth`` ``>=5.2.0`` ``gcloud-aio-bigquery`` ``>=6.1.2`` ``gcloud-aio-storage`` ``>=9.0.0`` ``gcsfs`` ``>=2023.10.0`` diff --git a/generated/provider_dependencies.json b/generated/provider_dependencies.json index ed1e4ab7b8283..58fb34b2adf02 100644 --- a/generated/provider_dependencies.json +++ b/generated/provider_dependencies.json @@ -605,7 +605,7 @@ "apache-airflow>=2.7.0", "asgiref>=3.5.2", "dill>=0.2.3", - "gcloud-aio-auth>=4.0.0,<5.0.0", + "gcloud-aio-auth>=5.2.0", "gcloud-aio-bigquery>=6.1.2", "gcloud-aio-storage>=9.0.0", "gcsfs>=2023.10.0", diff --git a/scripts/docker/entrypoint_ci.sh b/scripts/docker/entrypoint_ci.sh index 6d68e271d4851..1ff3ef0cd2cce 100755 --- a/scripts/docker/entrypoint_ci.sh +++ b/scripts/docker/entrypoint_ci.sh @@ -255,7 +255,7 @@ function check_boto_upgrade() { set -x # shellcheck disable=SC2086 ${PACKAGING_TOOL_CMD} install ${EXTRA_INSTALL_FLAGS} --upgrade boto3 botocore \ - "oss2>=2.14.0" "gcloud-aio-auth>=4.0.0,<5.0.0" "requests!=2.32.*,<3.0.0,>=2.24.0" + "oss2>=2.14.0" "cryptography<43.0.0" "requests!=2.32.*,<3.0.0,>=2.24.0" set +x pip check }