This is an automated email from the ASF dual-hosted git repository.

uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 011cd3debb Use offset-naive datetime in _CredentialsToken (#37539)
011cd3debb is described below

commit 011cd3debb4bb166908277c764d65eaf5985c7af
Author: Cedrik Neumann <7921017+m1rac...@users.noreply.github.com>
AuthorDate: Tue Feb 20 06:39:54 2024 +0100

    Use offset-naive datetime in _CredentialsToken (#37539)
    
    Co-authored-by: Andrey Anshin <andrey.ans...@taragol.is>
---
 .../providers/google/common/hooks/base_google.py    |  5 ++++-
 airflow/providers/google/provider.yaml              | 21 +++++++++++++--------
 .../google/common/hooks/test_base_google.py         |  4 ++++
 3 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/airflow/providers/google/common/hooks/base_google.py 
b/airflow/providers/google/common/hooks/base_google.py
index bee2798267..f45caeba5d 100644
--- a/airflow/providers/google/common/hooks/base_google.py
+++ b/airflow/providers/google/common/hooks/base_google.py
@@ -674,7 +674,10 @@ class _CredentialsToken(Token):
 
         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
index 317dc0e941..7edc6e137c 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 <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
index fd53930c5c..c251dbdb05 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 @@ class TestCredentialsToken:
         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"))

Reply via email to