Fix caching after migration to Niquests

Author: amotl

grafana_wtf/__init__.py

@@ -2,6 +2,3 @@
 
 __appname__ = "grafana-wtf"
 __version__ = "0.18.0"
-
-# Amalgamate `requests` to `niquests`.
-import grafana_wtf.compat

grafana_wtf/commands.py

@@ -234,7 +234,6 @@ def run():
 
     engine.enable_cache(expire_after=cache_ttl, drop_cache=options["drop-cache"])
     engine.enable_concurrency(int(options["concurrency"]))
-    engine.setup()
 
     log.info(f"Grafana version: {engine.version}")
 

grafana_wtf/compat.py

@@ -1,12 +1,10 @@
-from sys import modules
-
 import niquests
-import urllib3
+import requests_cache
+
+
+class CachedSession(requests_cache.session.CacheMixin, niquests.Session):
+    """
+    Make Niquests compatible with Requests-Cache.
+    """
 
-# Amalgamate the module namespace to make all modules aiming
-# to use `requests`, in fact use `niquests` instead.
-modules["requests"] = niquests
-modules["requests.adapters"] = niquests.adapters
-modules["requests.sessions"] = niquests.sessions
-modules["requests.exceptions"] = niquests.exceptions
-modules["requests.packages.urllib3"] = urllib3
+    pass

grafana_wtf/core.py

@@ -16,12 +16,12 @@
 from grafana_client.api import GrafanaApi
 from grafana_client.client import GrafanaClientError, GrafanaUnauthorizedError
 from munch import Munch, munchify
-from requests_cache import CachedSession
 from tqdm import tqdm
 from tqdm.contrib.logging import tqdm_logging_redirect
 from urllib3.exceptions import InsecureRequestWarning
 
 from grafana_wtf import __appname__, __version__
+from grafana_wtf.compat import CachedSession
 from grafana_wtf.model import (
     DashboardDetails,
     DashboardExplorationItem,
@@ -35,23 +35,34 @@
 
 
 class GrafanaEngine:
-    session = niquests.Session()
+
+    # Configure a larger HTTP request pool.
+    # TODO: Review the pool settings and eventually adjust according to concurrency level or other parameters.
+    # https://urllib3.readthedocs.io/en/latest/advanced-usage.html#customizing-pool-behavior
+    # https://laike9m.com/blog/requests-secret-pool_connections-and-pool_maxsize,89/
+    session_args = dict(pool_connections=100, pool_maxsize=100, retries=5)
+
+    # The HTTP `User-Agent` header value.
+    user_agent = f"{__appname__}/{__version__}"
 
     def __init__(self, grafana_url, grafana_token=None):
         self.grafana_url = grafana_url
         self.grafana_token = grafana_token
 
-        self.grafana = None
-        self.data = GrafanaDataModel()
+        self.concurrency = 5
 
+        self.grafana = self.grafana_client_factory(self.grafana_url, grafana_token=self.grafana_token)
+        self.set_user_agent()
+        self.data = GrafanaDataModel()
         self.finder = JsonPathFinder()
 
         self.taqadum = None
-        self.concurrency = 5
-
         self.debug = log.getEffectiveLevel() == logging.DEBUG
         self.progressbar = not self.debug
 
+    def set_session(self, session):
+        self.grafana.client.s = session
+
     def enable_cache(self, expire_after=60, drop_cache=False):
         if expire_after is None:
             log.info(f"Response cache will never expire (infinite caching)")
@@ -60,10 +71,13 @@ def enable_cache(self, expire_after=60, drop_cache=False):
         else:
             log.info(f"Response cache will expire after {expire_after} seconds")
 
-        self.session = CachedSession(cache_name=__appname__, expire_after=expire_after, use_cache_dir=True)
+        session = CachedSession(
+            cache_name=__appname__, expire_after=expire_after, use_cache_dir=True, **self.session_args
+        )
+        self.set_session(session)
+        self.set_user_agent()
 
-        cache = self.session.cache
-        log.info(f"Response cache database location is: {cache.db_path}")
+        log.info(f"Response cache database: {session.cache.db_path}")
         if drop_cache:
             log.info("Dropping response cache")
             self.clear_cache()
@@ -102,23 +116,14 @@ def grafana_client_factory(cls, grafana_url, grafana_token=None):
             url_path_prefix=url.path.lstrip("/"),
             verify=verify,
         )
-        if cls.session:
-            user_agent = f"{__appname__}/{__version__}"
-            cls.session.headers["User-Agent"] = user_agent
-            grafana.client.s = cls.session
 
-        return grafana
+        # Configure HTTP session to use a larger HTTP request pool.
+        grafana.client.s = niquests.Session(**cls.session_args)
 
-    def setup(self):
-        self.grafana = self.grafana_client_factory(self.grafana_url, grafana_token=self.grafana_token)
-
-        # Configure a larger HTTP request pool.
-        # Todo: Review the pool settings and eventually adjust according to concurrency level or other parameters.
-        # https://urllib3.readthedocs.io/en/latest/advanced-usage.html#customizing-pool-behavior
-        # https://laike9m.com/blog/requests-secret-pool_connections-and-pool_maxsize,89/
-        self.grafana.client.s = niquests.Session(pool_connections=100, pool_maxsize=100, retries=5)
+        return grafana
 
-        return self
+    def set_user_agent(self):
+        self.grafana.client.s.headers["User-Agent"] = self.user_agent
 
     def start_progressbar(self, total):
         if self.progressbar:

tests/conftest.py

@@ -320,7 +320,6 @@ def grafana_version(docker_grafana):
     Return Grafana version number.
     """
     engine = GrafanaWtf(grafana_url=docker_grafana, grafana_token=None)
-    engine.setup()
     grafana_version = engine.version
     return grafana_version
 

tests/test_core.py

@@ -40,7 +40,6 @@ def test_collect_datasource_items_variable_all():
 
 def test_connect_success():
     wtf = GrafanaWtf("https://play.grafana.org")
-    wtf.setup()
     health = wtf.health
     assert "commit" in health
     assert "version" in health
@@ -49,7 +48,6 @@ def test_connect_success():
 
 def test_connect_failure():
     wtf = GrafanaWtf("http://localhost:1234")
-    wtf.setup()
     with pytest.raises(ConnectionError) as ex:
         _ = wtf.health
     assert ex.match("The request to http://localhost:1234/api/health failed")
@@ -60,13 +58,11 @@ def test_connect_version(mock_get):
     mock_get.return_value = Mock()
     mock_get.return_value.return_value = {"commit": "14e988bd22", "database": "ok", "version": "9.0.1"}
     wtf = GrafanaWtf("http://localhost:1234")
-    wtf.setup()
     assert wtf.version == "9.0.1"
 
 
 def test_connect_non_json_response():
     wtf = GrafanaWtf("https://example.org/")
-    wtf.setup()
     with pytest.raises(ConnectionError) as ex:
         _ = wtf.health
     assert ex.match("The request to https://example.org/api/health failed: Client Error 404")