From ae4824e9e60ffb46d3d4d638c87b6da0fa7b5ac7 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Wed, 29 Jun 2022 18:32:12 +0530 Subject: [PATCH] Use token to validate the connection --- .../providers/google/common/hooks/base_google.py | 13 +++++++++++-- .../google/common/hooks/test_base_google.py | 8 ++++++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index aeacaa2a07089..4c4ab5c888c6f 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -31,6 +31,7 @@ import google.auth.credentials import google.oauth2.service_account import google_auth_httplib2 +import requests import tenacity from google.api_core.exceptions import Forbidden, ResourceExhausted, TooManyRequests from google.api_core.gapic_v1.client_info import ClientInfo @@ -270,7 +271,12 @@ def _get_credentials(self) -> google.auth.credentials.Credentials: def _get_access_token(self) -> str: """Returns a valid access token from Google API Credentials""" - return self._get_credentials().token + credentials = self._get_credentials() + auth_req = google.auth.transport.requests.Request() + # credentials.token is None + # Need to refresh credentials to populate the token + credentials.refresh(auth_req) + return credentials.token @functools.lru_cache(maxsize=None) def _get_credentials_email(self) -> str: @@ -585,7 +591,10 @@ def test_connection(self): """Test the Google cloud connectivity from UI""" status, message = False, '' try: - if self.project_id: + token = self._get_access_token() + url = f"https://www.googleapis.com/oauth2/v3/tokeninfo?access_token={token}" + response = requests.post(url) + if response.status_code == 200: status = True message = 'Connection successfully tested' except Exception as e: diff --git a/tests/providers/google/common/hooks/test_base_google.py b/tests/providers/google/common/hooks/test_base_google.py index 05187342e011a..247dce40b538c 100644 --- a/tests/providers/google/common/hooks/test_base_google.py +++ b/tests/providers/google/common/hooks/test_base_google.py @@ -341,9 +341,13 @@ def test_get_credentials_and_project_id_with_default_auth(self, mock_get_creds_a ) assert ('CREDENTIALS', 'PROJECT_ID') == result + @mock.patch('requests.post') @mock.patch(MODULE_NAME + '.get_credentials_and_project_id') - def test_connection_success(self, mock_get_creds_and_proj_id): - mock_get_creds_and_proj_id.return_value = ("CREDENTIALS", "PROJECT_ID") + def test_connection_success(self, mock_get_creds_and_proj_id, requests_post): + requests_post.return_value.status_code = 200 + credentials = mock.MagicMock() + type(credentials).token = mock.PropertyMock(return_value="TOKEN") + mock_get_creds_and_proj_id.return_value = (credentials, "PROJECT_ID") self.instance.extras = {} result = self.instance.test_connection() assert result == (True, 'Connection successfully tested')