From 73236e231a1cb7e3105a81f37997fb626ae6646d Mon Sep 17 00:00:00 2001 From: Hemslo Wang Date: Fri, 1 Sep 2023 04:01:08 +1000 Subject: [PATCH] Close connections after use (#2724) --- README.md | 3 ++ doc/introduction.rst | 4 ++ github/GithubIntegration.py | 17 ++++++++ github/MainClass.py | 17 ++++++++ github/Requester.py | 81 +++++++++++++++++++++++++++++-------- tests/Pickle.py | 26 ++++++++++++ tests/Requester.py | 28 +++++++++++++ 7 files changed, 160 insertions(+), 16 deletions(-) create mode 100644 tests/Pickle.py diff --git a/README.md b/README.md index 5de60200..2d82b2ee 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,9 @@ g = Github(base_url="https://{hostname}/api/v3", auth=auth) # Then play with your Github objects: for repo in g.get_user().get_repos(): print(repo.name) + +# To close connections after use +g.close() ``` ## Documentation diff --git a/doc/introduction.rst b/doc/introduction.rst index 88bf4fe0..18d6b04e 100644 --- a/doc/introduction.rst +++ b/doc/introduction.rst @@ -35,6 +35,10 @@ Then play with your Github objects:: # to see all the available attributes and methods print(dir(repo)) +To close connections after use:: + + g.close() + Download and install -------------------- diff --git a/github/GithubIntegration.py b/github/GithubIntegration.py index babba784..19d88510 100644 --- a/github/GithubIntegration.py +++ b/github/GithubIntegration.py @@ -119,6 +119,23 @@ class GithubIntegration: seconds_between_writes=seconds_between_writes, ) + def close(self) -> None: + """ + Close connections to the server. Alternatively, use the GithubIntegration object as a context manager: + + .. code-block:: python + + with github.GithubIntegration(...) as gi: + # do something + """ + self.__requester.close() + + def __enter__(self) -> "GithubIntegration": + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.close() + def get_github_for_installation(self, installation_id, token_permissions=None): # The installation has to authenticate as an installation, not an app auth = self.auth.get_installation_auth(installation_id, token_permissions, self.__requester) diff --git a/github/MainClass.py b/github/MainClass.py index d3010c2a..57d329e9 100644 --- a/github/MainClass.py +++ b/github/MainClass.py @@ -184,6 +184,23 @@ class Github: seconds_between_writes, ) + def close(self) -> None: + """ + Close connections to the server. Alternatively, use the Github object as a context manager: + + .. code-block:: python + + with github.Github(...) as gh: + # do something + """ + self.__requester.close() + + def __enter__(self) -> "Github": + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.close() + @property def FIX_REPO_GET_GIT_REF(self): """ diff --git a/github/Requester.py b/github/Requester.py index 62ce5951..6e710463 100644 --- a/github/Requester.py +++ b/github/Requester.py @@ -57,13 +57,28 @@ import logging import mimetypes import os import re +import threading import time import urllib import urllib.parse -from collections import defaultdict +from collections import deque from datetime import datetime, timezone from io import IOBase -from typing import TYPE_CHECKING, Any, Callable, Dict, Generic, ItemsView, List, Optional, Tuple, Type, TypeVar, Union +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Deque, + Dict, + Generic, + ItemsView, + List, + Optional, + Tuple, + Type, + TypeVar, + Union, +) import requests import requests.adapters @@ -166,7 +181,7 @@ class HTTPSRequestsConnectionClass: return RequestsResponse(r) def close(self) -> None: - return + self.session.close() class HTTPRequestsConnectionClass: @@ -229,7 +244,7 @@ class HTTPRequestsConnectionClass: return RequestsResponse(r) def close(self) -> None: - return + self.session.close() class Requester: @@ -238,7 +253,6 @@ class Requester: __httpConnectionClass = HTTPRequestsConnectionClass __httpsConnectionClass = HTTPSRequestsConnectionClass - __connection = None __persist = True __logger: Optional[logging.Logger] = None @@ -337,7 +351,6 @@ class Requester: __connectionClass: Union[Type[HTTPRequestsConnectionClass], Type[HTTPSRequestsConnectionClass]] __hostname: str __authorizationHeader: Optional[str] - __last_requests: Dict[str, float] __seconds_between_requests: Optional[float] __seconds_between_writes: Optional[float] @@ -369,7 +382,7 @@ class Requester: self.__pool_size = pool_size self.__seconds_between_requests = seconds_between_requests self.__seconds_between_writes = seconds_between_writes - self.__last_requests = defaultdict(lambda: 0.0) + self.__last_requests: Dict[str, float] = dict() self.__scheme = o.scheme if o.scheme == "https": self.__connectionClass = self.__httpsConnectionClass @@ -377,6 +390,9 @@ class Requester: self.__connectionClass = self.__httpConnectionClass else: assert False, "Unknown URL scheme" + self.__connection: Optional[Union[HTTPRequestsConnectionClass, HTTPSRequestsConnectionClass]] = None + self.__connection_lock = threading.Lock() + self.__custom_connections: Deque[Union[HTTPRequestsConnectionClass, HTTPSRequestsConnectionClass]] = deque() self.rate_limiting = (-1, -1) self.rate_limiting_resettime = 0 self.FIX_REPO_GET_GIT_REF = True @@ -397,6 +413,33 @@ class Requester: if isinstance(self.__auth, WithRequester): self.__auth.withRequester(self) + def __getstate__(self) -> Dict[str, Any]: + state = self.__dict__.copy() + # __connection_lock is not picklable + del state["_Requester__connection_lock"] + # __connection is not usable on remote, so ignore it + del state["_Requester__connection"] + # __custom_connections is not usable on remote, so ignore it + del state["_Requester__custom_connections"] + return state + + def __setstate__(self, state: Dict[str, Any]) -> None: + self.__dict__.update(state) + self.__connection_lock = threading.Lock() + self.__connection = None + self.__custom_connections = deque() + + def close(self) -> None: + """ + Close the connection to the server. + """ + with self.__connection_lock: + if self.__connection is not None: + self.__connection.close() + self.__connection = None + while self.__custom_connections: + self.__custom_connections.popleft().close() + @property def kwargs(self) -> Dict[str, Any]: """ @@ -499,6 +542,7 @@ class Requester: retry=self.__retry, pool_size=self.__pool_size, ) + self.__custom_connections.append(cnx) elif o.scheme == "https": cnx = self.__httpsConnectionClass( o.hostname, # type: ignore @@ -506,6 +550,7 @@ class Requester: retry=self.__retry, pool_size=self.__pool_size, ) + self.__custom_connections.append(cnx) return cnx @classmethod @@ -717,7 +762,6 @@ class Requester: responseHeaders = {k.lower(): v for k, v in response.getheaders()} output = response.read() - cnx.close() if input: if isinstance(input, IOBase): input.close() @@ -822,14 +866,19 @@ class Requester: if self.__persist and self.__connection is not None: return self.__connection - self.__connection = self.__connectionClass( - self.__hostname, - self.__port, - retry=self.__retry, - pool_size=self.__pool_size, - timeout=self.__timeout, - verify=self.__verify, - ) + with self.__connection_lock: + if self.__connection is not None: + if self.__persist: + return self.__connection + self.__connection.close() + self.__connection = self.__connectionClass( + self.__hostname, + self.__port, + retry=self.__retry, + pool_size=self.__pool_size, + timeout=self.__timeout, + verify=self.__verify, + ) return self.__connection diff --git a/tests/Pickle.py b/tests/Pickle.py new file mode 100644 index 00000000..3ab2e2ea --- /dev/null +++ b/tests/Pickle.py @@ -0,0 +1,26 @@ +import pickle +import unittest + +import github +from github.Repository import Repository + +REPO_NAME = "PyGithub/PyGithub" + + +class Pickle(unittest.TestCase): + def testPickleGithub(self): + gh = github.Github() + gh2 = pickle.loads(pickle.dumps(gh)) + self.assertIsInstance(gh2, github.Github) + self.assertIsNotNone(gh2._Github__requester._Requester__connection_lock) + self.assertIsNone(gh2._Github__requester._Requester__connection) + self.assertEqual(len(gh2._Github__requester._Requester__custom_connections), 0) + + def testPickleRepository(self): + gh = github.Github() + repo = gh.get_repo(REPO_NAME, lazy=True) + repo2 = pickle.loads(pickle.dumps(repo)) + self.assertIsInstance(repo2, Repository) + self.assertIsNotNone(repo2._requester._Requester__connection_lock) + self.assertIsNone(repo2._requester._Requester__connection) + self.assertEqual(len(repo2._requester._Requester__custom_connections), 0) diff --git a/tests/Requester.py b/tests/Requester.py index 6c6dda6f..85f825e6 100644 --- a/tests/Requester.py +++ b/tests/Requester.py @@ -26,6 +26,7 @@ from unittest import mock import github from . import Framework +from .GithubIntegration import APP_ID, PRIVATE_KEY REPO_NAME = "PyGithub/PyGithub" @@ -132,6 +133,33 @@ class Requester(Framework.TestCase): ), ) + def testCloseGithub(self): + mocked_connection = mock.MagicMock() + mocked_custom_connection = mock.MagicMock() + + with github.Github() as gh: + requester = gh._Github__requester + requester._Requester__connection = mocked_connection + requester._Requester__custom_connections.append(mocked_custom_connection) + + mocked_connection.close.assert_called_once_with() + mocked_custom_connection.close.assert_called_once_with() + self.assertIsNone(requester._Requester__connection) + + def testCloseGithubIntegration(self): + mocked_connection = mock.MagicMock() + mocked_custom_connection = mock.MagicMock() + + auth = github.Auth.AppAuth(APP_ID, PRIVATE_KEY) + with github.GithubIntegration(auth=auth) as gi: + requester = gi._GithubIntegration__requester + requester._Requester__connection = mocked_connection + requester._Requester__custom_connections.append(mocked_custom_connection) + + mocked_connection.close.assert_called_once_with() + mocked_custom_connection.close.assert_called_once_with() + self.assertIsNone(requester._Requester__connection) + def testLoggingRedirection(self): self.assertEqual(self.g.get_repo("EnricoMi/test").name, "test-renamed") self.logger.info.assert_called_once_with(