From 9011edadafececcd46d2b9eac7e5b4930ff35bd3 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Mon, 27 Jun 2022 21:56:18 +0530 Subject: [PATCH 1/2] Add test_connection method to GoogleBaseHook --- .../providers/google/common/hooks/base_google.py | 13 +++++++++++++ .../google/common/hooks/test_base_google.py | 14 ++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index 997c72c4e6ac7..aeacaa2a07089 100644 --- a/airflow/providers/google/common/hooks/base_google.py +++ b/airflow/providers/google/common/hooks/base_google.py @@ -580,3 +580,16 @@ def download_content_from_request(file_handle, request: dict, chunk_size: int) - while done is False: _, done = downloader.next_chunk() file_handle.flush() + + def test_connection(self): + """Test the Google cloud connectivity from UI""" + status, message = False, '' + try: + if self.project_id: + status = True + message = 'Connection successfully tested' + except Exception as e: + status = False + message = str(e) + + return status, message diff --git a/tests/providers/google/common/hooks/test_base_google.py b/tests/providers/google/common/hooks/test_base_google.py index a60d3a4e22840..05187342e011a 100644 --- a/tests/providers/google/common/hooks/test_base_google.py +++ b/tests/providers/google/common/hooks/test_base_google.py @@ -341,6 +341,20 @@ def test_get_credentials_and_project_id_with_default_auth(self, mock_get_creds_a ) assert ('CREDENTIALS', 'PROJECT_ID') == result + @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") + self.instance.extras = {} + result = self.instance.test_connection() + assert result == (True, 'Connection successfully tested') + + @mock.patch(MODULE_NAME + '.get_credentials_and_project_id') + def test_connection_failure(self, mock_get_creds_and_proj_id): + mock_get_creds_and_proj_id.side_effect = AirflowException('Invalid key JSON.') + self.instance.extras = {} + result = self.instance.test_connection() + assert result == (False, 'Invalid key JSON.') + @mock.patch(MODULE_NAME + '.get_credentials_and_project_id') def test_get_credentials_and_project_id_with_service_account_file(self, mock_get_creds_and_proj_id): mock_credentials = mock.MagicMock() From ae4824e9e60ffb46d3d4d638c87b6da0fa7b5ac7 Mon Sep 17 00:00:00 2001 From: Phani Kumar Date: Wed, 29 Jun 2022 18:32:12 +0530 Subject: [PATCH 2/2] 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')