Update GoogleBaseHook to not follow 308 and use 60s timeout (#8816)
This commit is contained in:
Родитель
7d69987edd
Коммит
e1e833bb26
|
@ -32,14 +32,13 @@ import google.auth
|
|||
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
|
||||
|
@ -213,7 +212,7 @@ class GoogleBaseHook(BaseHook):
|
|||
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
|
||||
|
|
|
@ -591,7 +591,7 @@ class TestGoogleBaseHook(unittest.TestCase):
|
|||
}
|
||||
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):
|
||||
"""
|
||||
|
@ -616,6 +616,22 @@ class TestGoogleBaseHook(unittest.TestCase):
|
|||
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):
|
||||
|
|
Загрузка…
Ссылка в новой задаче