diff --git a/hexlib/db.py b/hexlib/db.py index 0fce394..075cb09 100644 --- a/hexlib/db.py +++ b/hexlib/db.py @@ -3,11 +3,11 @@ import sqlite3 import traceback import psycopg2 -import redis -import orjson as json import umsgpack from psycopg2.errorcodes import UNIQUE_VIOLATION +from hexlib.env import get_redis + class PersistentState: """Quick and dirty persistent dict-like SQLite wrapper""" @@ -25,8 +25,8 @@ class PersistentState: class VolatileState: """Quick and dirty volatile dict-like redis wrapper""" - def __init__(self, prefix, **redis_args): - self.rdb = redis.Redis(**redis_args) + def __init__(self, prefix, redis_db): + self.rdb = redis_db self.prefix = prefix def __getitem__(self, table): @@ -36,8 +36,10 @@ class VolatileState: class VolatileQueue: """Quick and dirty volatile queue-like redis wrapper""" - def __init__(self, key, **redis_args): - self.rdb = redis.Redis(**redis_args) + def __init__(self, key, redis_db=None): + if redis_db is None: + redis_db = get_redis() + self.rdb = redis_db self.key = key def put(self, item): @@ -52,8 +54,8 @@ class VolatileQueue: class VolatileBooleanState: """Quick and dirty volatile dict-like redis wrapper for boolean values""" - def __init__(self, prefix, **redis_args): - self.rdb = redis.Redis(**redis_args) + def __init__(self, prefix, redis_db): + self.rdb = redis_db self.prefix = prefix def __getitem__(self, table): @@ -184,7 +186,6 @@ class Table: pass - def _sqlite_type(value): if isinstance(value, int): return "integer" diff --git a/hexlib/env.py b/hexlib/env.py new file mode 100644 index 0000000..8085b87 --- /dev/null +++ b/hexlib/env.py @@ -0,0 +1,22 @@ +import redis +import os + +from hexlib.log import stdout_logger +from hexlib.web import Web + + +def get_redis(): + return redis.Redis( + host=os.environ.get("REDIS_HOST", "localhost"), + port=int(os.environ.get("REDIS_PORT", 6379)) + ) + + +def get_web(): + return Web( + proxy=os.environ.get("PROXY", None), + rps=os.environ.get("RPS", 1), + logger=stdout_logger, + cookie_file=os.environ.get("COOKIE_FILE", None), + retry_codes=set(os.environ.get("RETRY_CODES", "").split(",")) + ) diff --git a/hexlib/log.py b/hexlib/log.py new file mode 100644 index 0000000..0803267 --- /dev/null +++ b/hexlib/log.py @@ -0,0 +1,51 @@ +import logging +import os +import sys +from logging import StreamHandler + +DATE_FMT = "%Y-%m-%d %H:%M:%S" + + +class ColorFormatter(logging.Formatter): + + def __init__(self, fmt): + super().__init__() + + grey = "\x1b[38;21m" + yellow = "\x1b[33;21m" + red = "\x1b[31;21m" + bold_red = "\x1b[31;1m" + reset = "\x1b[0m" + + self.formats = { + logging.DEBUG: logging.Formatter(grey + fmt + reset, datefmt=DATE_FMT), + logging.INFO: logging.Formatter(grey + fmt + reset, datefmt=DATE_FMT), + logging.WARNING: logging.Formatter(yellow + fmt + reset, datefmt=DATE_FMT), + logging.ERROR: logging.Formatter(red + fmt + reset, datefmt=DATE_FMT), + logging.CRITICAL: logging.Formatter(bold_red + fmt + reset, datefmt=DATE_FMT) + } + + def format(self, record): + return self.formats[record.levelno].format(record) + + +stdout_logger = logging.getLogger("default") + +if os.environ.get("LOG_LEVEL", "debug") == "debug": + stdout_logger.setLevel(logging.DEBUG) + +for h in stdout_logger.handlers: + stdout_logger.removeHandler(h) + +handler = StreamHandler(sys.stdout) +if os.environ.get("LOG_THREAD_NAME", "0") == "1": + fmt = "%(asctime)s %(levelname)-5s>%(threadName)s %(message)s" +else: + fmt = "%(asctime)s %(levelname)-5s>%(message)s" + +if os.environ.get("LOG_COLORS", "1") == "1": + handler.formatter = ColorFormatter(fmt) +else: + handler.formatter = logging.Formatter(fmt, datefmt='%Y-%m-%d %H:%M:%S') +stdout_logger.addHandler(handler) +logger = stdout_logger diff --git a/hexlib/web.py b/hexlib/web.py index 74cc9cc..202b01b 100644 --- a/hexlib/web.py +++ b/hexlib/web.py @@ -4,12 +4,16 @@ import os from datetime import datetime from base64 import b64encode, b64decode from http.cookiejar import Cookie +from time import time + import requests import orjson as json from dateutil.parser import parse from requests.cookies import RequestsCookieJar +from hexlib.misc import rate_limit, retry + def cookie_from_string(text: str, domain: str) -> Cookie: tokens = [t.strip() for t in text.split(";")] @@ -104,3 +108,63 @@ def download_file(url, destination, session=None, headers=None, overwrite=False, if err_cb: err_cb(e) retries -= 1 + +class Web: + def __init__(self, proxy=None, rps=1, retries=3, logger=None, cookie_file=None, retry_codes=None, session=None): + self._cookie_file = cookie_file + self._proxy = proxy + self._logger = logger + self._current_req = None + if retry_codes is None: + retry_codes = {502, 504, 522, 524, 429} + self._retry_codes = retry_codes + + if session is None: + session = requests.session() + + self._session = session + + if self._cookie_file: + self._session.cookies = load_cookiejar(cookie_file) + + if self._proxy: + self._session.proxies = { + "http": proxy, + "https": proxy, + } + + @rate_limit(rps) + @retry(retries, callback=self._error_callback) + def get(url, **kwargs): + self._current_req = "GET", url, kwargs + r = self._session.get(url, **kwargs) + + if r.status_code in self._retry_codes: + raise Exception(f"HTTP {r.status_code}") + return r + + self._get = get + + def _error_callback(self, e): + self._logger.critical(f"{self._format_url(*self._current_req)}: {e}") + + def _format_url(self, method, url, kwargs, r=None): + if "params" in kwargs and kwargs["params"]: + return "%s %s?%s <%s>" % (method, url, "&".join(f"{k}={v}" for k, v in kwargs["params"].items()), + r.status_code if r else "ERR") + else: + return "%s %s <%s>" % (method, url, r.status_code if r else "ERR",) + + def get(self, url, **kwargs): + + time_start = time() + r = self._get(url, **kwargs) + + if self._cookie_file: + save_cookiejar(self._session.cookies, self._cookie_file) + + if self._logger and r is not None: + self._logger.debug(self._format_url("GET", url, kwargs, r) + " %.2fs" % (time() - time_start)) + return r + + diff --git a/setup.py b/setup.py index d81c970..0f6dc50 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ from setuptools import setup setup( name="hexlib", - version="1.28", + version="1.29", description="Misc utility methods", author="simon987", author_email="me@simon987.net", diff --git a/test/test_VolatileState.py b/test/test_VolatileState.py index deabce9..15ba3f4 100644 --- a/test/test_VolatileState.py +++ b/test/test_VolatileState.py @@ -1,4 +1,5 @@ from unittest import TestCase + from hexlib.db import VolatileState, VolatileBooleanState, VolatileQueue diff --git a/test/test_download_file.py b/test/test_download_file.py index 52ab797..28f4fc7 100644 --- a/test/test_download_file.py +++ b/test/test_download_file.py @@ -1,5 +1,5 @@ -from unittest import TestCase import os +from unittest import TestCase from hexlib.web import download_file diff --git a/test/test_retry.py b/test/test_retry.py index e4896b8..67c7ca3 100644 --- a/test/test_retry.py +++ b/test/test_retry.py @@ -1,5 +1,5 @@ from unittest import TestCase -from hexlib.db import VolatileState, VolatileBooleanState + from hexlib.misc import retry