Skip to content

Commit

Permalink
Update GoogleBaseHook to not follow 308 and use 60s timeout (#8816)
Browse files Browse the repository at this point in the history
  • Loading branch information
waiyan1612 committed May 13, 2020
1 parent 7d69987 commit e1e833b
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 4 deletions.
5 changes: 2 additions & 3 deletions airflow/providers/google/common/hooks/base_google.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,13 @@
import google.auth.credentials
import google.oauth2.service_account
import google_auth_httplib2
import httplib2
import tenacity
from google.api_core.exceptions import Forbidden, ResourceExhausted, TooManyRequests
from google.api_core.gapic_v1.client_info import ClientInfo
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 @@ -213,7 +212,7 @@ def _authorize(self) -> google_auth_httplib2.AuthorizedHttp:
service hook connection.
"""
credentials = self._get_credentials()
http = httplib2.Http()
http = build_http()
http = set_user_agent(http, "airflow/" + version.version)
authed_http = google_auth_httplib2.AuthorizedHttp(credentials, http=http)
return authed_http
Expand Down
18 changes: 17 additions & 1 deletion tests/providers/google/common/hooks/test_base_google.py
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.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,22 @@ def test_authorize_assert_user_agent_is_sent(self, mock_get_credentials, mock_ht
self.assertEqual(response, new_response)
self.assertEqual(content, new_content)

@mock.patch("airflow.providers.google.common.hooks.base_google.GoogleBaseHook._get_credentials")
def test_authorize_assert_http_308_is_excluded(self, mock_get_credentials):
"""
Verify that 308 status code is excluded from httplib2's redirect codes
"""
http_authorized = self.instance._authorize().http
self.assertTrue(308 not in http_authorized.redirect_codes)

@mock.patch("airflow.providers.google.common.hooks.base_google.GoogleBaseHook._get_credentials")
def test_authorize_assert_http_timeout_is_present(self, mock_get_credentials):
"""
Verify that http client has a timeout set
"""
http_authorized = self.instance._authorize().http
self.assertNotEqual(http_authorized.timeout, None)


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

0 comments on commit e1e833b

Please sign in to comment.