Close connections after use (#2724)

hemslo · web-flow · commit 73236e231a1c · 2023-08-31T20:01:08.000+02:00
diff --git 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
@@ -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
@@ -119,6 +119,23 @@ def __init__(
             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
@@ -184,6 +184,23 @@ def __init__(
             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
@@ -57,13 +57,28 @@
 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 @@ def getresponse(self) -> RequestsResponse:
         return RequestsResponse(r)
 
     def close(self) -> None:
-        return
+        self.session.close()
 
 
 class HTTPRequestsConnectionClass:
@@ -229,7 +244,7 @@ def getresponse(self) -> RequestsResponse:
         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 @@ def _initializeDebugFeature(self) -> None:
     __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,14 +382,17 @@ def __init__(
         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
         elif o.scheme == "http":
             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 @@ def __init__(
         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,13 +542,15 @@ def __customConnection(
                         retry=self.__retry,
                         pool_size=self.__pool_size,
                     )
+                    self.__custom_connections.append(cnx)
                 elif o.scheme == "https":
                     cnx = self.__httpsConnectionClass(
                         o.hostname,  # type: ignore
                         o.port,
                         retry=self.__retry,
                         pool_size=self.__pool_size,
                     )
+                    self.__custom_connections.append(cnx)
         return cnx
 
     @classmethod
@@ -717,7 +762,6 @@ def __requestRaw(
             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 @@ def __createConnection(
         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
@@ -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
@@ -26,6 +26,7 @@
 import github
 
 from . import Framework
+from .GithubIntegration import APP_ID, PRIVATE_KEY
 
 REPO_NAME = "PyGithub/PyGithub"
 
@@ -132,6 +133,33 @@ class TestAuth(github.Auth.AppAuth):
             ),
         )
 
+    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(
