add session_pool_size

grafana-toolbox · Oct 18, 2024 · 36a2386 · 36a2386
1 parent 49226b0
commit 36a2386
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/grafana_client/api.py b/grafana_client/api.py
@@ -8,7 +8,7 @@
 import niquests.auth
 from urllib3.exceptions import InsecureRequestWarning
 
-from .client import DEFAULT_TIMEOUT, AsyncGrafanaClient, GrafanaClient
+from .client import DEFAULT_SESSION_POOL_SIZE, DEFAULT_TIMEOUT, AsyncGrafanaClient, GrafanaClient
 from .elements import (
     Admin,
     Alerting,
@@ -72,6 +72,7 @@ def __init__(
         timeout=DEFAULT_TIMEOUT,
         user_agent: str = None,
         organization_id: int = None,
+        session_pool_size=DEFAULT_SESSION_POOL_SIZE,
     ):
         self.client = GrafanaClient(
             auth,
@@ -83,6 +84,7 @@ def __init__(
             timeout=timeout,
             user_agent=user_agent,
             organization_id=organization_id,
+            session_pool_size=session_pool_size,
         )
         self.url = None
         self.admin = Admin(self.client)

diff --git a/grafana_client/client.py b/grafana_client/client.py
@@ -88,6 +88,7 @@ def __init__(
         timeout=DEFAULT_TIMEOUT,
         user_agent: str = None,
         organization_id: int = None,
+        session_pool_size=DEFAULT_SESSION_POOL_SIZE,
     ):
         self.auth = auth
         self.verify = verify
@@ -96,6 +97,7 @@ def __init__(
         self.url_port = port
         self.url_path_prefix = url_path_prefix
         self.url_protocol = protocol
+        self.session_pool_size = session_pool_size
 
         def construct_api_url():
             params = {
@@ -118,7 +120,7 @@ def construct_api_url():
 
         self.user_agent = user_agent or f"{__appname__}/{__version__}"
 
-        self.s = niquests.Session()
+        self.s = niquests.Session(pool_maxsize=session_pool_size)
         self.s.headers["User-Agent"] = self.user_agent
 
         self.organization_id = organization_id
@@ -230,6 +232,7 @@ def __init__(
         timeout=DEFAULT_TIMEOUT,
         user_agent: str = None,
         organization_id: int = None,
+        session_pool_size=DEFAULT_SESSION_POOL_SIZE,
     ):
         super().__init__(
             auth,
@@ -241,8 +244,9 @@ def __init__(
             timeout=timeout,
             user_agent=user_agent,
             organization_id=organization_id,
+            session_pool_size=session_pool_size,
         )
-        self.s = niquests.AsyncSession()
+        self.s = niquests.AsyncSession(pool_maxsize=session_pool_size)
         self.s.headers.setdefault("Connection", "keep-alive")
 
     def __getattr__(self, item):