sf-veritas 0.9.7__py3-none-any.whl

sf_veritas/patches/network_libraries/niquests.py

@@ -0,0 +1,192 @@
+"""
+Monkey-patch Niquests so that every flavour of request
+(sync / async, Session / AsyncSession, streaming or not)
+propagates the SAILFISH_TRACING_HEADER *and* records the request.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import List, Optional, Tuple
+
+from ...constants import SAILFISH_TRACING_HEADER
+from ..constants import supported_network_verbs as verbs
+from .utils import get_trace_and_should_propagate, record_network_request
+
+
+def patch_niquests(domains_to_not_propagate_headers_to: Optional[List[str]] = None):
+    try:
+        import niquests  # type: ignore
+    except ImportError:
+        return
+
+    skip = domains_to_not_propagate_headers_to or []
+
+    # --------------------------------------------------------------------- helpers
+    SessionCls = niquests.Session
+    AsyncSessionCls = getattr(niquests, "AsyncSession", type("._Dummy", (), {}))
+
+    def _is_session_self(args: Tuple) -> bool:
+        """True when args[0] is a Session/AsyncSession instance."""
+        return bool(args) and isinstance(args[0], (SessionCls, AsyncSessionCls))
+
+    def _resolve_method_url(
+        verb_tag: str, args: Tuple, kwargs: dict
+    ) -> Tuple[str, str]:
+        """
+        Robust extraction of (method, url) for every call style.
+        Never throws IndexError – falls back to kwargs when positional args missing.
+        """
+        # ----- 1. REQUEST / AREQUEST (generic entry-points) ----------------
+        if verb_tag in ("REQUEST", "AREQUEST"):
+            if _is_session_self(args):
+                method = (
+                    str(args[1]).upper()
+                    if len(args) > 1
+                    else str(kwargs.get("method", "GET")).upper()
+                )
+                url = args[2] if len(args) > 2 else kwargs.get("url", "")
+            else:
+                method = (
+                    str(args[0]).upper()
+                    if len(args) > 0
+                    else str(kwargs.get("method", "GET")).upper()
+                )
+                url = args[1] if len(args) > 1 else kwargs.get("url", "")
+
+        # ----- 2. Convenience verbs get/post/… and their async variants ----
+        else:
+            method = verb_tag.lstrip("A").upper()
+            if _is_session_self(args):
+                url = args[1] if len(args) > 1 else kwargs.get("url", "")
+            else:
+                url = args[0] if len(args) > 0 else kwargs.get("url", "")
+
+        return method, url
+
+    def _prepare(url: str, method: str, headers: Optional[dict]):
+        trace_id, allow = get_trace_and_should_propagate(url, skip)
+        hdrs = dict(headers or {})
+        if allow:
+            hdrs[SAILFISH_TRACING_HEADER] = trace_id
+        return trace_id, hdrs, int(time.time() * 1_000)
+
+    def _record(trace_id, url, method, status, success, err, start_ms):
+        record_network_request(
+            trace_id,
+            url,
+            method,
+            status or 0,
+            success,
+            err,
+            timestamp_start=start_ms,
+            timestamp_end=int(time.time() * 1_000),
+        )
+
+    # ----------------------------------------------------------------- header merge
+    def _headers_pos_index(args: Tuple) -> Optional[int]:
+        """
+        Return the positional index that already holds a headers dict, or None.
+        """
+        if not args:
+            return None
+        if isinstance(args[0], (SessionCls, AsyncSessionCls)):
+            return 3 if len(args) > 3 else None
+        return 2 if len(args) > 2 else None
+
+    def _inject_header(args: Tuple, kwargs: dict, key: str, val: str) -> Tuple:
+        idx = _headers_pos_index(args)
+        if idx is not None:
+            merged = dict(args[idx] or {})
+            merged[key] = val
+            args = (*args[:idx], merged, *args[idx + 1 :])
+        else:
+            hdrs = dict(kwargs.get("headers") or {})
+            hdrs[key] = val
+            kwargs["headers"] = hdrs
+        return args, kwargs
+
+    # -------------------------------------------------------------------- wrappers
+    def _wrap_sync(fn, verb_tag: str):
+        def wrapper(*args, **kwargs):
+            method, url = _resolve_method_url(verb_tag, args, kwargs)
+            trace_id, _hdrs, t0 = _prepare(url, method, kwargs.get("headers"))
+            args, kwargs = _inject_header(
+                args, kwargs, SAILFISH_TRACING_HEADER, trace_id
+            )
+
+            status = 0
+            success = False
+            err = None
+            try:
+                resp = fn(*args, **kwargs)
+                status = getattr(resp, "status_code", 0)
+                success = True
+                return resp
+            except Exception as exc:  # noqa: BLE001
+                err = str(exc)[:255]
+                raise
+            finally:
+                _record(trace_id, url, method, status, success, err, t0)
+
+        return wrapper
+
+    def _wrap_async(fn, verb_tag: str):
+        async def wrapper(*args, **kwargs):
+            method, url = _resolve_method_url(verb_tag, args, kwargs)
+            trace_id, _hdrs, t0 = _prepare(url, method, kwargs.get("headers"))
+            args, kwargs = _inject_header(
+                args, kwargs, SAILFISH_TRACING_HEADER, trace_id
+            )
+
+            status = 0
+            success = False
+            err = None
+            try:
+                resp = await fn(*args, **kwargs)
+                status = getattr(resp, "status_code", 0)
+                success = True
+                return resp
+            except Exception as exc:  # noqa: BLE001
+                err = str(exc)[:255]
+                raise
+            finally:
+                _record(trace_id, url, method, status, success, err, t0)
+
+        return wrapper
+
+    # ------------------------------------------------------------- apply patches
+    niquests.request = _wrap_sync(niquests.request, "REQUEST")
+    for v in verbs:
+        setattr(niquests, v, _wrap_sync(getattr(niquests, v), v.upper()))
+
+    SessionCls.request = _wrap_sync(SessionCls.request, "REQUEST")
+    for v in verbs:
+        setattr(SessionCls, v, _wrap_sync(getattr(SessionCls, v), v.upper()))
+
+    if hasattr(niquests, "arequest"):
+        niquests.arequest = _wrap_async(niquests.arequest, "AREQUEST")
+        for av in verbs:
+            async_name = f"a{av}"
+            setattr(
+                niquests,
+                async_name,
+                _wrap_async(getattr(niquests, async_name), async_name.upper()),
+            )
+
+    if hasattr(SessionCls, "arequest"):
+        SessionCls.arequest = _wrap_async(SessionCls.arequest, "AREQUEST")
+        for av in verbs:
+            async_name = f"a{av}"
+            setattr(
+                SessionCls,
+                async_name,
+                _wrap_async(getattr(SessionCls, async_name), async_name.upper()),
+            )
+
+    if AsyncSessionCls is not type("._Dummy", (), {}):
+        AsyncSessionCls.request = _wrap_async(AsyncSessionCls.request, "REQUEST")
+        for v in verbs:
+            setattr(
+                AsyncSessionCls, v, _wrap_async(getattr(AsyncSessionCls, v), v.upper())
+            )

sf_veritas/patches/network_libraries/pycurl.py

@@ -0,0 +1,71 @@
+import time
+from typing import List, Optional
+
+from ...constants import SAILFISH_TRACING_HEADER
+from .utils import get_trace_and_should_propagate, record_network_request
+
+
+def patch_pycurl(domains_to_not_propagate_headers_to: Optional[List[str]] = None):
+    try:
+        import pycurl
+    except ImportError:
+        return
+
+    _OrigCurl = pycurl.Curl
+
+    class WrappedCurl(_OrigCurl):  # ➊ subclass libcurl handle
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            self._sf_url: str | None = None
+            self._sf_method: str | None = None
+            self._sf_headers: list[str] = []
+
+        # --- intercept option setting -------------------------------------------------
+        def setopt(self, opt, val):
+            if opt == pycurl.URL:
+                self._sf_url = val
+            elif opt == pycurl.CUSTOMREQUEST:
+                self._sf_method = val.upper()
+            elif opt == pycurl.HTTPHEADER:
+                self._sf_headers = list(val)
+            return super().setopt(opt, val)
+
+        # --- wrapped perform() --------------------------------------------------------
+        def perform(self):
+            url = self._sf_url or ""
+            method = (self._sf_method or "GET").upper()
+
+            trace_id, allow = get_trace_and_should_propagate(
+                url, domains_to_not_propagate_headers_to or []
+            )
+
+            # Build merged header list
+            merged = list(self._sf_headers)
+            if allow:
+                merged.append(f"{SAILFISH_TRACING_HEADER}: {trace_id}")
+
+            # Let libcurl negotiate & decode encodings for us
+            super().setopt(pycurl.ACCEPT_ENCODING, "")
+
+            # push merged headers down
+            # NOTE: HTTPHEADER expects List[str] (or List[bytes]), ensure consistency
+            super().setopt(pycurl.HTTPHEADER, merged)
+
+            # timing / status / error capture
+            ts0 = int(time.time() * 1_000)
+            status = 0
+            err: str | None = None
+            try:
+                rv = super().perform()
+                status = int(self.getinfo(pycurl.RESPONSE_CODE) or 0)
+                return rv
+            except Exception as e:
+                err = str(e)[:255]
+                raise
+            finally:
+                ts1 = int(time.time() * 1_000)
+                record_network_request(
+                    trace_id, url, method, status, err is None, err, ts0, ts1
+                )
+
+    pycurl.Curl = WrappedCurl

sf_veritas/patches/network_libraries/requests.py

@@ -0,0 +1,187 @@
+"""
+Monkey-patch the `requests` stack (requests → urllib3 → http.client):
+
+• For every outbound request, propagate the SAILFISH_TRACING_HEADER
+  unless the destination host is in `domains_to_not_propagate_headers_to`.
+• Fire NetworkRequestTransmitter via utils.record_network_request
+  so we always capture (url, status, timings, success, error).
+"""
+
+from __future__ import annotations
+
+import http.client
+import time
+from typing import Dict, List, Optional, Tuple
+
+import requests
+import urllib3
+from requests.sessions import Session
+
+from ...constants import SAILFISH_TRACING_HEADER
+from ...thread_local import (
+    activate_reentrancy_guards_exception,
+    activate_reentrancy_guards_logging,
+    activate_reentrancy_guards_print,
+)
+from .utils import get_trace_and_should_propagate, record_network_request
+
+###############################################################################
+# Internal helpers
+###############################################################################
+
+# header names used for re-entrancy guards
+REENTRANCY_GUARD_LOGGING_PREACTIVE = "reentrancy_guard_logging_preactive"
+REENTRANCY_GUARD_PRINT_PREACTIVE = "reentrancy_guard_print_preactive"
+REENTRANCY_GUARD_EXCEPTIONS_PREACTIVE = "reentrancy_guard_exception_preactive"
+
+
+def _activate_rg(headers: Dict[str, str]) -> None:
+    """Turn the three ‘preactive' guard flags ON for downstream hops."""
+    headers[REENTRANCY_GUARD_LOGGING_PREACTIVE] = "true"
+    headers[REENTRANCY_GUARD_PRINT_PREACTIVE] = "true"
+    headers[REENTRANCY_GUARD_EXCEPTIONS_PREACTIVE] = "true"
+
+
+def _check_rg(headers: Dict[str, str]) -> None:
+    """If any pre-active guard present, switch the corresponding guard on."""
+    if headers.get(REENTRANCY_GUARD_LOGGING_PREACTIVE, "false").lower() == "true":
+        activate_reentrancy_guards_logging()
+    if headers.get(REENTRANCY_GUARD_PRINT_PREACTIVE, "false").lower() == "true":
+        activate_reentrancy_guards_print()
+    if headers.get(REENTRANCY_GUARD_EXCEPTIONS_PREACTIVE, "false").lower() == "true":
+        activate_reentrancy_guards_exception()
+
+
+def _prepare(
+    url: str,
+    domains_to_skip: List[str],
+    headers: Optional[Dict[str, str]],
+) -> Tuple[str, Dict[str, str], int]:
+    """
+    Inject the trace header (unless excluded) and return:
+        trace_id, merged_headers, timestamp_ms
+    """
+    trace_id, propagate = get_trace_and_should_propagate(url, domains_to_skip)
+    hdrs: Dict[str, str] = dict(headers or {})
+    _check_rg(hdrs)
+    if propagate:
+        hdrs[SAILFISH_TRACING_HEADER] = trace_id
+    _activate_rg(hdrs)
+    return trace_id, hdrs, int(time.time() * 1_000)
+
+
+###############################################################################
+# Top-level patch function
+###############################################################################
+def patch_requests(domains_to_not_propagate_headers_to: Optional[List[str]] = None):
+    """Apply all monkey-patches.  Safe to call multiple times."""
+    exclude = domains_to_not_propagate_headers_to or []
+
+    # --------------------------------------------------------------------- #
+    # 1. Patch `requests.Session.request`
+    # --------------------------------------------------------------------- #
+    original_request = Session.request
+
+    def patched_request(self: Session, method, url, **kwargs):  # type: ignore[override]
+        # --- header handling / injection --------------------------------- #
+        trace_id, hdrs, t0 = _prepare(url, exclude, kwargs.pop("headers", {}))
+        kwargs["headers"] = hdrs
+
+        status: int = 0
+        success: bool = False
+        err: str | None = None
+        try:
+            resp = original_request(self, method, url, **kwargs)
+            status = resp.status_code
+            success = resp.ok
+            return resp
+        except Exception as exc:  # noqa: BLE001
+            err = str(exc)[:255]
+            raise
+        finally:
+            record_network_request(
+                trace_id,
+                url,
+                str(method).upper(),
+                status,
+                success,
+                err,
+                timestamp_start=t0,
+                timestamp_end=int(time.time() * 1_000),
+            )
+
+    Session.request = patched_request
+    requests.Session.request = patched_request  # cover direct `requests.Session(...)`
+
+    # --------------------------------------------------------------------- #
+    # 2. Patch urllib3's low-level ConnectionPool.urlopen (used by requests)
+    # --------------------------------------------------------------------- #
+    original_urlopen = urllib3.connectionpool.HTTPConnectionPool.urlopen
+
+    def patched_urlopen(self, method, url, body=None, headers=None, **kw):  # type: ignore[override]
+        trace_id, hdrs, t0 = _prepare(url, exclude, headers)
+        status: int = 0
+        success: bool = False
+        err: str | None = None
+        try:
+            resp = original_urlopen(self, method, url, body=body, headers=hdrs, **kw)
+            status = getattr(resp, "status", 0)
+            success = status < 400
+            return resp
+        except Exception as exc:  # noqa: BLE001
+            err = str(exc)[:255]
+            raise
+        finally:
+            record_network_request(
+                trace_id,
+                url,
+                str(method).upper(),
+                status,
+                success,
+                err,
+                timestamp_start=t0,
+                timestamp_end=int(time.time() * 1_000),
+            )
+
+    urllib3.connectionpool.HTTPConnectionPool.urlopen = patched_urlopen
+
+    # --------------------------------------------------------------------- #
+    # 3. Patch http.client for “raw” stdlib usage (rare but easy to support)
+    # --------------------------------------------------------------------- #
+    original_http_client_request = http.client.HTTPConnection.request
+
+    def patched_http_request(self, method, url, body=None, headers=None, *, encode_chunked=False):  # type: ignore[override]
+        trace_id, hdrs, t0 = _prepare(url, exclude, headers)
+        status: int = 0
+        success: bool = False
+        err: str | None = None
+        try:
+            resp = original_http_client_request(
+                self,
+                method,
+                url,
+                body=body,
+                headers=hdrs,
+                encode_chunked=encode_chunked,
+            )
+            status = getattr(
+                self, "response", getattr(resp, "status", 0)
+            )  # best-effort
+            success = bool(status) and status < 400
+            return resp
+        except Exception as exc:  # noqa: BLE001
+            err = str(exc)[:255]
+            raise
+        finally:
+            record_network_request(
+                trace_id,
+                url,
+                str(method).upper(),
+                status,
+                success,
+                err,
+                timestamp_start=t0,
+                timestamp_end=int(time.time() * 1_000),
+            )
+
+    http.client.HTTPConnection.request = patched_http_request

sf_veritas/patches/network_libraries/tornado.py

@@ -0,0 +1,139 @@
+"""
+Monkey-patch Tornado's HTTP clients so that
+
+• Every outbound request carries SAILFISH_TRACING_HEADER
+  (unless the destination host is excluded).
+• Every request – success or failure – triggers record_network_request(…).
+
+Covers
+  • tornado.httpclient.AsyncHTTPClient.fetch   (await-able)
+  • tornado.httpclient.HTTPClient.fetch        (blocking/sync)
+Safe to call repeatedly; patches only once per process.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import List, Optional, Tuple
+
+from ...constants import SAILFISH_TRACING_HEADER
+from .utils import get_trace_and_should_propagate, record_network_request
+
+
+def patch_tornado(domains_to_not_propagate_headers_to: Optional[List[str]] = None):
+    try:
+        # Tornado is optional; exit silently if missing
+        from tornado.httpclient import AsyncHTTPClient, HTTPClient, HTTPRequest
+    except ImportError:
+        return
+
+    exclude: List[str] = domains_to_not_propagate_headers_to or []
+
+    # ------------------------------------------------------------------ #
+    # Helpers shared by sync & async wrappers
+    # ------------------------------------------------------------------ #
+    def _resolve(
+        req_or_url, kwargs
+    ) -> Tuple[str, str, dict]:  # → (url, METHOD, headers_dict)
+        """
+        Handle both call styles:
+
+            client.fetch("https://foo", method="POST", headers={...})
+            client.fetch(HTTPRequest(...))
+
+        Always returns a mutable *headers* dict.
+        """
+        if isinstance(req_or_url, HTTPRequest):
+            url = req_or_url.url
+            method = (req_or_url.method or "GET").upper()
+            hdrs = dict(req_or_url.headers or {})
+        else:
+            url = str(req_or_url)
+            method = kwargs.get("method", "GET").upper()
+            hdrs = dict(kwargs.get("headers", {}) or {})
+        return url, method, hdrs
+
+    def _inject(
+        req_or_url, kwargs, hdrs: dict
+    ):  # mutate request object *or* kwargs to carry hdrs
+        from tornado.httpclient import HTTPRequest  # local import to avoid MRO issues
+
+        if isinstance(req_or_url, HTTPRequest):
+            req_or_url.headers = hdrs
+        else:
+            kwargs["headers"] = hdrs
+        return req_or_url, kwargs
+
+    def _prepare(url: str, hdrs: dict):
+        """Return (trace_id, merged_headers, start_ms)."""
+        trace_id, allow = get_trace_and_should_propagate(url, exclude)
+        out = dict(hdrs)
+        if allow:
+            out[SAILFISH_TRACING_HEADER] = trace_id
+        return trace_id, out, int(time.time() * 1_000)
+
+    # ------------------------------------------------------------------ #
+    # AsyncHTTPClient.fetch wrapper
+    # ------------------------------------------------------------------ #
+    original_async_fetch = AsyncHTTPClient.fetch
+
+    async def patched_async_fetch(self, req_or_url, *args, **kwargs):
+        url, method, hdrs_cur = _resolve(req_or_url, kwargs)
+        trace_id, hdrs_new, t0 = _prepare(url, hdrs_cur)
+        req_or_url, kwargs = _inject(req_or_url, kwargs, hdrs_new)
+
+        status, success, err = 0, False, None
+        try:
+            resp = await original_async_fetch(self, req_or_url, *args, **kwargs)
+            status = getattr(resp, "code", 0)
+            success = status < 400
+            return resp
+        except Exception as exc:  # noqa: BLE001
+            err = str(exc)[:255]
+            raise
+        finally:
+            record_network_request(
+                trace_id,
+                url,
+                method,
+                status,
+                success,
+                err,
+                timestamp_start=t0,
+                timestamp_end=int(time.time() * 1_000),
+            )
+
+    AsyncHTTPClient.fetch = patched_async_fetch  # type: ignore[assignment]
+
+    # ------------------------------------------------------------------ #
+    # HTTPClient.fetch wrapper (blocking)
+    # ------------------------------------------------------------------ #
+    original_sync_fetch = HTTPClient.fetch
+
+    def patched_sync_fetch(self, req_or_url, *args, **kwargs):
+        url, method, hdrs_cur = _resolve(req_or_url, kwargs)
+        trace_id, hdrs_new, t0 = _prepare(url, hdrs_cur)
+        req_or_url, kwargs = _inject(req_or_url, kwargs, hdrs_new)
+
+        status, success, err = 0, False, None
+        try:
+            resp = original_sync_fetch(self, req_or_url, *args, **kwargs)
+            status = getattr(resp, "code", 0)
+            success = status < 400
+            return resp
+        except Exception as exc:  # noqa: BLE001
+            err = str(exc)[:255]
+            raise
+        finally:
+            record_network_request(
+                trace_id,
+                url,
+                method,
+                status,
+                success,
+                err,
+                timestamp_start=t0,
+                timestamp_end=int(time.time() * 1_000),
+            )
+
+    HTTPClient.fetch = patched_sync_fetch  # type: ignore[assignment]