From 98888dadf04e7e00524b6de273d28d02d7abc2c0 Mon Sep 17 00:00:00 2001 From: Bu Sun Kim <8822365+busunkim96@users.noreply.github.com> Date: Wed, 23 Sep 2020 11:10:39 -0600 Subject: [PATCH] fix: add method to close httplib2 connections (#1038) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #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 --- docs/start.md | 16 ++++++++++++++-- googleapiclient/discovery.py | 29 +++++++++++++++++++++++++++-- googleapiclient/http.py | 2 ++ tests/test_discovery.py | 26 ++++++++++++++++++++++++++ 4 files changed, 69 insertions(+), 4 deletions(-) diff --git a/docs/start.md b/docs/start.md index 2db95faed11..ceac5dd3364 100644 --- 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 index eec7e003446..6363809f1e5 100644 --- 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 index d9c3d2ae19a..926ca1bcfa9 100644 --- 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 index 1abb5c82f08..c6bc59937b8 100644 --- 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/"