fix: add method to close httplib2 connections (googleapis#1038)

busunkim96 · web-flow · commit 98888dadf04e · 2020-09-23T11:10:39.000-06:00
Fixes googleapis#618 🦕 httplib2 leaves sockets open by default. This can lead to situations where programs run out of available sockets. httplib2 added a method last year to clean up connections. https://github.com/httplib2/httplib2/blob/9bf300cdc372938f4237150d5b9b615879eb51a1/python3/httplib2/__init__.py#L1498-L1506 This PR adds two ways to close http connections. The interface is intentionally similar to the proposed design for GAPIC clients. googleapis/gapic-generator-python#575
diff --git a/docs/start.md b/docs/start.md
@@ -45,12 +45,24 @@ This section describes how to build an API-specific service object, make calls t
 
 ### Build the service object
 
-Whether you are using simple or authorized API access, you use the [build()](http://googleapis.github.io/google-api-python-client/docs/epy/googleapiclient.discovery-module.html#build) function to create a service object. It takes an API name and API version as arguments. You can see the list of all API versions on the [Supported APIs](dyn/index.md) page. The service object is constructed with methods specific to the given API. To create it, do the following:
+Whether you are using simple or authorized API access, you use the [build()](http://googleapis.github.io/google-api-python-client/docs/epy/googleapiclient.discovery-module.html#build) function to create a service object. It takes an API name and API version as arguments. You can see the list of all API versions on the [Supported APIs](dyn/index.md) page. The service object is constructed with methods specific to the given API. 
+
+`httplib2`, the underlying transport library, makes all connections persistent by default. Use the service object with a context manager or call `close` to avoid leaving sockets open.
 
 
 ```python
 from googleapiclient.discovery import build
-service = build('api_name', 'api_version', ...)
+
+service = build('drive', 'v3')
+# ...
+service.close()
+```
+
+```python
+from googleapiclient.discovery import build
+
+with build('drive', 'v3') as service:
+    # ...
 ```
 
 ### Collections
diff --git a/googleapiclient/discovery.py b/googleapiclient/discovery.py
@@ -261,6 +261,8 @@ def build(
     else:
         discovery_http = http
 
+    service = None
+
     for discovery_url in _discovery_service_uri_options(discoveryServiceUrl, version):
         requested_url = uritemplate.expand(discovery_url, params)
 
@@ -273,7 +275,7 @@ def build(
                 developerKey,
                 num_retries=num_retries,
             )
-            return build_from_document(
+            service = build_from_document(
                 content,
                 base=discovery_url,
                 http=http,
@@ -285,13 +287,22 @@ def build(
                 adc_cert_path=adc_cert_path,
                 adc_key_path=adc_key_path,
             )
+            break  # exit if a service was created
         except HttpError as e:
             if e.resp.status == http_client.NOT_FOUND:
                 continue
             else:
                 raise e
 
-    raise UnknownApiNameOrVersion("name: %s  version: %s" % (serviceName, version))
+    # If discovery_http was created by this function, we are done with it
+    # and can safely close it
+    if http is None:
+        discovery_http.close()
+
+    if service is None:
+        raise UnknownApiNameOrVersion("name: %s  version: %s" % (serviceName, version))
+    else:
+        return service
 
 
 def _discovery_service_uri_options(discoveryServiceUrl, version):
@@ -1309,6 +1320,20 @@ def __setstate__(self, state):
         self._dynamic_attrs = []
         self._set_service_methods()
 
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, exc_tb):
+        self.close()
+
+    def close(self):
+        """Close httplib2 connections."""
+        # httplib2 leaves sockets open by default.
+        # Cleanup using the `close` method.
+        # https://github.com/httplib2/httplib2/issues/148
+        self._http.http.close()
+
     def _set_service_methods(self):
         self._add_basic_methods(self._resourceDesc, self._rootDesc, self._schema)
         self._add_nested_resources(self._resourceDesc, self._rootDesc, self._schema)
diff --git a/googleapiclient/http.py b/googleapiclient/http.py
@@ -1720,6 +1720,8 @@ def request(
         self.headers = headers
         return httplib2.Response(self.response_headers), self.data
 
+    def close(self):
+        return None
 
 class HttpMockSequence(object):
     """Mock of httplib2.Http
diff --git a/tests/test_discovery.py b/tests/test_discovery.py
@@ -437,6 +437,13 @@ def test_ResourceMethodParameters_zoo_animals_patch(self):
         self.assertEqual(parameters.enum_params, {})
 
 
+class Discovery(unittest.TestCase):
+    def test_discovery_http_is_closed(self):
+        http = HttpMock(datafile("malformed.json"), {"status": "200"})
+        service = build("plus", "v1", credentials=mock.sentinel.credentials)
+        http.close.assert_called_once()
+
+
 class DiscoveryErrors(unittest.TestCase):
     def test_tests_should_be_run_with_strict_positional_enforcement(self):
         try:
@@ -572,6 +579,25 @@ def test_building_with_developer_key_skips_adc(self):
         # application default credentials were used.
         self.assertNotIsInstance(plus._http, google_auth_httplib2.AuthorizedHttp)
 
+    def test_building_with_context_manager(self):
+        discovery = read_datafile("plus.json")
+        with mock.patch("httplib2.Http") as http:
+            with build_from_document(discovery, base="https://www.googleapis.com/", credentials=self.MOCK_CREDENTIALS) as plus:
+                self.assertIsNotNone(plus)
+                self.assertTrue(hasattr(plus, "activities"))
+            plus._http.http.close.assert_called_once()
+
+    def test_resource_close(self):
+        discovery = read_datafile("plus.json")
+        with mock.patch("httplib2.Http") as http:
+            plus = build_from_document(
+                discovery,
+                base="https://www.googleapis.com/",
+                credentials=self.MOCK_CREDENTIALS,
+            )
+            plus.close()
+            plus._http.http.close.assert_called_once()
+
     def test_api_endpoint_override_from_client_options(self):
         discovery = read_datafile("plus.json")
         api_endpoint = "https://foo.googleapis.com/"
