Use offset-naive datetime in _CredentialsToken (#37539)

Co-authored-by: Andrey Anshin <[email protected]>
apache · Feb 20, 2024 · 011cd3d · 011cd3d
1 parent 7b422be
commit 011cd3d
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 9 deletions.
diff --git 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.now(tz=datetime.timezone.utc).replace(tzinfo=None)
  self.acquiring = None
 
 

diff --git 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 <https://ads.google.com/>`__
-  - `Google Cloud (GCP) <https://cloud.google.com/>`__
-  - `Google Firebase <https://firebase.google.com/>`__
-  - `Google LevelDB <https://github.com/google/leveldb/>`__
-  - `Google Marketing Platform <https://marketingplatform.google.com/>`__
-  - `Google Workspace <https://workspace.google.com/>`__ (formerly Google Suite)
+ - `Google Ads <https://ads.google.com/>`__
+ - `Google Cloud (GCP) <https://cloud.google.com/>`__
+ - `Google Firebase <https://firebase.google.com/>`__
+ - `Google LevelDB <https://github.com/google/leveldb/>`__
+ - `Google Marketing Platform <https://marketingplatform.google.com/>`__
+ - `Google Workspace <https://workspace.google.com/>`__ (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:

diff --git 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"))