diff --git a/CHANGES.rst b/CHANGES.rst index c53d8b0..283b3b2 100644 --- a/CHANGES.rst +++ b/CHANGES.rst @@ -12,6 +12,8 @@ in progress - Add subcommand ``explore permissions``. Thanks, @meyerder. - Added support for Python 3.12 - Removed support for Python 3.7 +- SQLite cache: Use ``requests_cache.CachedSession`` for better concurrency + behaviour. Thanks, @JensRichnow and @JWCook. 2024-03-07 0.18.0 ================= diff --git a/grafana_wtf/__init__.py b/grafana_wtf/__init__.py index fbd2318..47852bb 100644 --- a/grafana_wtf/__init__.py +++ b/grafana_wtf/__init__.py @@ -2,3 +2,6 @@ __appname__ = "grafana-wtf" __version__ = "0.18.0" + +# Amalgamate `requests` to `niquests`. +import grafana_wtf.compat diff --git a/grafana_wtf/compat.py b/grafana_wtf/compat.py new file mode 100644 index 0000000..57a8205 --- /dev/null +++ b/grafana_wtf/compat.py @@ -0,0 +1,12 @@ +from sys import modules + +import niquests +import urllib3 + +# 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 diff --git a/grafana_wtf/core.py b/grafana_wtf/core.py index 2fa1bd7..8ac0fba 100644 --- a/grafana_wtf/core.py +++ b/grafana_wtf/core.py @@ -16,6 +16,7 @@ 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 @@ -34,6 +35,8 @@ class GrafanaEngine: + session = niquests.Session() + def __init__(self, grafana_url, grafana_token): self.grafana_url = grafana_url self.grafana_token = grafana_token @@ -56,9 +59,11 @@ def enable_cache(self, expire_after=60, drop_cache=False): log.info(f"Response cache will expire immediately (expire_after=0)") else: log.info(f"Response cache will expire after {expire_after} seconds") - requests_cache.install_cache(cache_name=__appname__, expire_after=expire_after, use_cache_dir=True) - cache_database_file = requests_cache.get_cache().db_path - log.info(f"Response cache database location is {cache_database_file}") + + self.session = CachedSession(cache_name=__appname__, expire_after=expire_after, use_cache_dir=True) + + cache = self.session.cache + log.info(f"Response cache database location is: {cache.db_path}") if drop_cache: log.info("Dropping response cache") self.clear_cache() @@ -72,8 +77,8 @@ def clear_cache(self): def enable_concurrency(self, concurrency): self.concurrency = concurrency - @staticmethod - def grafana_client_factory(grafana_url, grafana_token=None): + @classmethod + def grafana_client_factory(cls, grafana_url, grafana_token=None): url = urlparse(grafana_url) # Grafana API Key auth @@ -97,6 +102,9 @@ def grafana_client_factory(grafana_url, grafana_token=None): url_path_prefix=url.path.lstrip("/"), verify=verify, ) + if cls.session: + cls.session.headers["User-Agent"] = grafana.client.user_agent + grafana.client.s = cls.session return grafana diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29