Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update GoogleBaseHook to not follow 308 and use 60s timeout #8816

Merged
merged 4 commits into from
May 13, 2020
Merged

Update GoogleBaseHook to not follow 308 and use 60s timeout #8816

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
d3f801a
Exclude 308 status code from httplib2's list of redirect codes.
waiyan1612 May 11, 2020
3eeac05
Use build_http from google-api-client.
waiyan1612 May 12, 2020
cf4b4d1
Fix tests.
waiyan1612 May 12, 2020
98a283f
Do not expose build_http method.
waiyan1612 May 12, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions airflow/providers/google/common/hooks/base_google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@
from google.auth import _cloud_sdk
from google.auth.environment_vars import CREDENTIALS
from googleapiclient.errors import HttpError
from googleapiclient.http import MediaIoBaseDownload, set_user_agent
from googleapiclient.http import MediaIoBaseDownload, build_http, set_user_agent

from airflow import version
from airflow.exceptions import AirflowException
Expand Down Expand Up @@ -207,13 +207,20 @@ def _get_access_token(self) -> str:
"""
return self._get_credentials().token

@staticmethod
def _build_http() -> httplib2.Http:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We do not need to give the opportunity to override this method. If the user needs it, the user can overwrite authorize. It would make sense if _http were used in several places.

"""
Returns a HTTP object using the defaults from google api client
"""
return build_http()

def _authorize(self) -> google_auth_httplib2.AuthorizedHttp:
"""
Returns an authorized HTTP object to be used to build a Google cloud
service hook connection.
"""
credentials = self._get_credentials()
http = httplib2.Http()
http = self._build_http()
http = set_user_agent(http, "airflow/" + version.version)
authed_http = google_auth_httplib2.AuthorizedHttp(credentials, http=http)
return authed_http
Expand Down
16 changes: 15 additions & 1 deletion tests/providers/google/common/hooks/test_base_google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -591,7 +591,7 @@ def test_num_retries_is_not_none_by_default(self, get_con_mock):
}
self.assertEqual(self.instance.num_retries, 5)

@mock.patch("airflow.providers.google.common.hooks.base_google.httplib2.Http")
@mock.patch("airflow.providers.google.common.hooks.base_google.GoogleBaseHook._build_http")
@mock.patch("airflow.providers.google.common.hooks.base_google.GoogleBaseHook._get_credentials")
def test_authorize_assert_user_agent_is_sent(self, mock_get_credentials, mock_http):
"""
Expand All @@ -616,6 +616,20 @@ def test_authorize_assert_user_agent_is_sent(self, mock_get_credentials, mock_ht
self.assertEqual(response, new_response)
self.assertEqual(content, new_content)

def test_build_http_308_is_excluded(self):
"""
Verify that 308 status code is excluded from httplib2's redirect codes
"""
http = self.instance._build_http()
self.assertTrue(308 not in http.redirect_codes)

def test_build_http_timeout_is_present(self):
"""
Verify that http client has a timeout set
"""
http = self.instance._build_http()
self.assertNotEqual(http.timeout, None)


class TestProvideAuthorizedGcloud(unittest.TestCase):
def setUp(self):
Expand Down