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

Author: waiyan1612

airflow/providers/google/common/hooks/base_google.py

@@ -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
@@ -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

tests/providers/google/common/hooks/test_base_google.py

@@ -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):
         """
@@ -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):