diff --git a/airflow/providers/google/common/hooks/base_google.py b/airflow/providers/google/common/hooks/base_google.py index 997c72c4e6ac71..4c4ab5c888c6f6 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: @@ -580,3 +586,19 @@ 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: + 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: + 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 a60d3a4e228405..247dce40b538cf 100644 --- a/tests/providers/google/common/hooks/test_base_google.py +++ b/tests/providers/google/common/hooks/test_base_google.py @@ -341,6 +341,24 @@ 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, 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') + + @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()