From 43caaf9db3a8d445f1b0b741540d51e8dffb4b09 Mon Sep 17 00:00:00 2001
From: Johannes Nussbaum <39048939+jnussbaum@users.noreply.github.com>
Date: Tue, 16 Jan 2024 09:59:00 +0100
Subject: [PATCH] chore: improve performance with requests.Session object
 (DEV-3174) (#739)

---
 src/dsp_tools/utils/connection_live.py | 44 ++++++++++----------------
 1 file changed, 17 insertions(+), 27 deletions(-)

diff --git a/src/dsp_tools/utils/connection_live.py b/src/dsp_tools/utils/connection_live.py
index bf2277807..ab9ecd4c8 100644
--- a/src/dsp_tools/utils/connection_live.py
+++ b/src/dsp_tools/utils/connection_live.py
@@ -7,8 +7,7 @@
 from typing import Any, Callable, Optional, cast
 
 import regex
-import requests
-from requests import ReadTimeout, RequestException, Response
+from requests import ReadTimeout, RequestException, Response, Session
 from urllib3.exceptions import ReadTimeoutError
 
 from dsp_tools.models.exceptions import BaseError
@@ -34,11 +33,15 @@ class ConnectionLive:
 
     server: str
     token: Optional[str] = None
+    session: Session = field(init=False, default=Session())
     # downtimes of server-side services -> API still processes request
     # -> retry too early has side effects (e.g. duplicated resources)
     timeout_put_post: int = field(init=False, default=30 * 60)
     timeout_get_delete: int = field(init=False, default=20)
 
+    def __post_init__(self) -> None:
+        self.session.headers["User-Agent"] = f'DSP-TOOLS/{version("dsp-tools")}'
+
     def login(self, email: str, password: str) -> None:
         """
         Retrieve a session token and store it as class attribute.
@@ -62,6 +65,7 @@ def login(self, email: str, password: str) -> None:
                 api_route="/v2/authentication",
             )
         self.token = response["token"]
+        self.session.headers["Authorization"] = f"Bearer {self.token}"
 
     def logout(self) -> None:
         """
@@ -102,7 +106,9 @@ def _log_request(
                 _return["token"] = "<token>"
         else:
             _return = {"status": response.status_code, "message": response.text}
-        if headers and "Authorization" in headers:
+        headers = headers or {}
+        headers.update({k: str(v) for k, v in self.session.headers.items()})
+        if "Authorization" in headers:
             headers["Authorization"] = regex.sub(r"Bearer .+", "Bearer <token>", headers["Authorization"])
         if data and "password" in data:
             data["password"] = "<password>"
@@ -146,16 +152,12 @@ def post(
         if not route.startswith("/"):
             route = f"/{route}"
         url = self.server + route
-        if not headers:
-            headers = {}
-        headers["User-Agent"] = f'DSP-TOOLS/{version("dsp-tools")}'
         if data:
+            headers = headers or {}
             headers["Content-Type"] = "application/json; charset=UTF-8"
-        if self.token:
-            headers["Authorization"] = f"Bearer {self.token}"
         timeout = timeout or self.timeout_put_post
 
-        request = partial(requests.post, url=url, headers=headers, timeout=timeout)
+        request = partial(self.session.post, url=url, headers=headers, timeout=timeout)
         if data:
             # if data is not encoded as bytes, issues can occur with non-ASCII characters,
             # where the content-length of the request will turn out to be different from the actual length
@@ -195,15 +197,10 @@ def get(
         if not route.startswith("/"):
             route = f"/{route}"
         url = self.server + route
-        if not headers:
-            headers = {}
-        headers["User-Agent"] = f'DSP-TOOLS/{version("dsp-tools")}'
-        if self.token:
-            headers["Authorization"] = f"Bearer {self.token}"
         timeout = self.timeout_get_delete
 
         response = self._try_network_action(
-            lambda: requests.get(
+            lambda: self.session.get(
                 url=url,
                 headers=headers,
                 timeout=timeout,
@@ -242,17 +239,13 @@ def put(
         if not route.startswith("/"):
             route = f"/{route}"
         url = self.server + route
-        if not headers:
-            headers = {}
-        headers["User-Agent"] = f'DSP-TOOLS/{version("dsp-tools")}'
         if data:
+            headers = headers or {}
             headers["Content-Type"] = f"{content_type}; charset=UTF-8"
-        if self.token:
-            headers["Authorization"] = f"Bearer {self.token}"
         timeout = self.timeout_put_post
 
         response = self._try_network_action(
-            lambda: requests.put(
+            lambda: self.session.put(
                 url=url,
                 headers=headers,
                 # if data is not encoded as bytes, issues can occur with non-ASCII characters,
@@ -292,14 +285,9 @@ def delete(
         if not route.startswith("/"):
             route = f"/{route}"
         url = self.server + route
-        if not headers:
-            headers = {}
-        headers["User-Agent"] = f'DSP-TOOLS/{version("dsp-tools")}'
-        if self.token:
-            headers["Authorization"] = f"Bearer {self.token}"
         timeout = self.timeout_get_delete
 
-        response = requests.delete(
+        response = self.session.delete(
             url=url,
             headers=headers,
             params=params,
@@ -352,6 +340,8 @@ def _try_network_action(self, action: Callable[[], Response]) -> Response:
                 self._log_and_sleep(reason="Timeout Error", retry_counter=i)
                 continue
             except (ConnectionError, RequestException):
+                self.session.close()
+                self.session = Session()
                 self._log_and_sleep(reason="Network Error", retry_counter=i)
                 continue