From 08ab9ecaa4ef8ebd252ccc19328790df68639378 Mon Sep 17 00:00:00 2001 From: Cedrik Neumann Date: Mon, 19 Feb 2024 11:49:00 +0100 Subject: [PATCH 1/3] fix: use offset-naive datetime in _CredentialsToken The field `access_token_acquired_at` is specific to gcloud-aio's Token class and will be used on subsequent calls of `get` with an offset-naive datetime insternally. Therefore we cannot use offset-aware instances of datetime here. --- airflow/providers/google/common/hooks/base_google.py | 5 ++++- tests/providers/google/common/hooks/test_base_google.py | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index bee2798267610..ec6297d0ac3f5 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -674,7 +674,10 @@ async def acquire_access_token(self, timeout: int = 10) -> None: self.access_token = cast(str, self.credentials.token) self.access_token_duration = 3600 - self.access_token_acquired_at = datetime.datetime.now(tz=datetime.timezone.utc) + # 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.utcnow() # noqa: TID251 self.acquiring = None diff --git a/tests/providers/google/common/hooks/test_base_google.py b/tests/providers/google/common/hooks/test_base_google.py index fd53930c5c1ff..c251dbdb05636 100644 --- a/tests/providers/google/common/hooks/test_base_google.py +++ b/tests/providers/google/common/hooks/test_base_google.py @@ -892,6 +892,10 @@ async def test_get(self): token = hook._CredentialsToken(mock_credentials, project=PROJECT_ID, scopes=SCOPES) assert await token.get() == "ACCESS_TOKEN" mock_credentials.refresh.assert_called_once() + # ensure token caching works on subsequent calls of `get` + mock_credentials.reset_mock() + assert await token.get() == "ACCESS_TOKEN" + mock_credentials.refresh.assert_not_called() @pytest.mark.asyncio @mock.patch(f"{MODULE_NAME}.get_credentials_and_project_id", return_value=("CREDENTIALS", "PROJECT_ID")) From 5d073673e1a1c28d6c580b9559d1a7fb1d0eacda Mon Sep 17 00:00:00 2001 From: Cedrik Neumann <7921017+m1racoli@users.noreply.github.com> Date: Mon, 19 Feb 2024 12:21:44 +0100 Subject: [PATCH 2/3] Update airflow/providers/google/common/hooks/base_google.py Co-authored-by: Andrey Anshin --- airflow/providers/google/common/hooks/base_google.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index ec6297d0ac3f5..f45caeba5dad4 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -677,7 +677,7 @@ async def acquire_access_token(self, timeout: int = 10) -> None: # 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.utcnow() # noqa: TID251 + self.access_token_acquired_at = datetime.datetime.now(tz=datetime.timezone.utc).replace(tzinfo=None) self.acquiring = None From fcf3da25d7da05fc287bfb97eeb565eaba140b5e Mon Sep 17 00:00:00 2001 From: Cedrik Neumann Date: Mon, 19 Feb 2024 12:35:52 +0100 Subject: [PATCH 3/3] docs: add hints for upgrading gcloud-aio-auth Let others know about things to consider when upgrading the major version of gcloud-aio-auth. --- airflow/providers/google/provider.yaml | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/airflow/providers/google/provider.yaml b/airflow/providers/google/provider.yaml index 317dc0e9410c8..7edc6e137c7c7 100644 --- a/airflow/providers/google/provider.yaml +++ b/airflow/providers/google/provider.yaml @@ -19,14 +19,14 @@ package-name: apache-airflow-providers-google name: Google description: | - Google services including: + Google services including: - - `Google Ads `__ - - `Google Cloud (GCP) `__ - - `Google Firebase `__ - - `Google LevelDB `__ - - `Google Marketing Platform `__ - - `Google Workspace `__ (formerly Google Suite) + - `Google Ads `__ + - `Google Cloud (GCP) `__ + - `Google Firebase `__ + - `Google LevelDB `__ + - `Google Marketing Platform `__ + - `Google Workspace `__ (formerly Google Suite) state: ready source-date-epoch: 1707636385 @@ -89,6 +89,12 @@ dependencies: - apache-airflow>=2.6.0 - apache-airflow-providers-common-sql>=1.7.2 - asgiref>=3.5.2 + # 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-bigquery>=6.1.2 - gcloud-aio-storage>=9.0.0 @@ -677,7 +683,6 @@ operators: python-modules: - airflow.providers.google.cloud.operators.cloud_batch - sensors: - integration-name: Google BigQuery python-modules: