logxpy 0.1.0__py3-none-any.whl

logxpy/dask.py

@@ -0,0 +1,172 @@
+"""Support for Eliot tracing with Dask computations."""
+
+from pyrsistent import PClass, field
+
+from dask import compute, optimize, persist
+
+try:
+    from dask.distributed import Future
+    from dask.highlevelgraph import HighLevelGraph
+except:
+
+    class Future(object):
+        pass
+
+
+from dask.core import toposort, get_dependencies, ishashable
+from . import start_action, current_action, Action
+
+
+class _RunWithEliotContext(PClass):
+    """
+    Run a callable within an Eliot context.
+
+    @ivar task_id: The serialized Eliot task ID.
+    @ivar func: The function that Dask wants to run.
+    @ivar key: The key in the Dask graph.
+    @ivar dependencies: The keys in the Dask graph this depends on.
+    """
+
+    task_id = field(type=str)
+    func = field()  # callable
+    key = field(type=str)
+    dependencies = field()
+
+    # Pretend to be underlying callable for purposes of equality; necessary for
+    # optimizer to be happy:
+
+    def __eq__(self, other):
+        return self.func == other
+
+    def __ne__(self, other):
+        return self.func != other
+
+    def __hash__(self):
+        return hash(self.func)
+
+    def __call__(self, *args, **kwargs):
+        with Action.continue_task(task_id=self.task_id) as action:
+            action.log(
+                message_type="dask:task", key=self.key, dependencies=self.dependencies
+            )
+            return self.func(*args, **kwargs)
+
+
+def compute_with_trace(*args):
+    """Do Dask compute(), but with added Eliot tracing.
+
+    Dask is a graph of tasks, but Eliot logs trees.  So we need to emulate a
+    graph using a tree.  We do this by making Eliot action for each task, but
+    having it list the tasks it depends on.
+
+    We use the following algorithm:
+
+        1. Create a top-level action.
+
+        2. For each entry in the dask graph, create a child with
+           serialize_task_id.  Do this in likely order of execution, so that
+           if B depends on A the task level of B is higher than the task Ievel
+           of A.
+
+        3. Replace each function with a wrapper that uses the corresponding
+           task ID (with Action.continue_task), and while it's at it also
+           records which other things this function depends on.
+
+    Known issues:
+
+        1. Retries will confuse Eliot.  Probably need different
+           distributed-tree mechanism within Eliot to solve that.
+    """
+    # 1. Create top-level Eliot Action:
+    with start_action(action_type="dask:compute"):
+        # In order to reduce logging verbosity, add logging to the already
+        # optimized graph:
+        optimized = optimize(*args, optimizations=[_add_logging])
+        return compute(*optimized, optimize_graph=False)
+
+
+def persist_with_trace(*args):
+    """Do Dask persist(), but with added Eliot tracing.
+
+    Known issues:
+
+        1. Retries will confuse Eliot.  Probably need different
+           distributed-tree mechanism within Eliot to solve that.
+    """
+    # 1. Create top-level Eliot Action:
+    with start_action(action_type="dask:persist"):
+        # In order to reduce logging verbosity, add logging to the already
+        # optimized graph:
+        optimized = optimize(*args, optimizations=[_add_logging])
+        return persist(*optimized, optimize_graph=False)
+
+
+def _add_logging(dsk, ignore=None):
+    """
+    Add logging to a Dask graph.
+
+    @param dsk: The Dask graph.
+
+    @return: New Dask graph.
+    """
+    if isinstance(dsk, HighLevelGraph):
+        dsk = dsk.to_dict()
+
+    ctx = current_action()
+    result = {}
+
+    # Use topological sort to ensure Eliot actions are in logical order of
+    # execution in Dask:
+    keys = toposort(dsk)
+
+    # Give each key a string name. Some keys are just aliases to other
+    # keys, so make sure we have underlying key available. Later on might
+    # want to shorten them as well.
+    def simplify(k):
+        if isinstance(k, str):
+            return k
+        return "-".join(str(o) for o in k)
+
+    key_names = {}
+    for key in keys:
+        value = dsk[key]
+        if not callable(value) and ishashable(value) and value in keys:
+            # It's an alias for another key:
+            key_names[key] = key_names[value]
+        else:
+            key_names[key] = simplify(key)
+
+    # Values in the graph can be either:
+    #
+    # 1. A list of other values.
+    # 2. A tuple, where first value might be a callable, aka a task.
+    # 3. A literal of some sort.
+    def maybe_wrap(key, value):
+        if isinstance(value, list):
+            return [maybe_wrap(key, v) for v in value]
+        elif isinstance(value, tuple):
+            func = value[0]
+            args = value[1:]
+            if not callable(func):
+                # Not a callable, so nothing to wrap.
+                return value
+            wrapped_func = _RunWithEliotContext(
+                task_id=str(ctx.serialize_task_id(), "utf-8"),
+                func=func,
+                key=key_names[key],
+                dependencies=[key_names[k] for k in get_dependencies(dsk, key)],
+            )
+            return (wrapped_func,) + args
+        else:
+            return value
+
+    # Replace function with wrapper that logs appropriate Action; iterate in
+    # topological order so action task levels are in reasonable order.
+    for key in keys:
+        result[key] = maybe_wrap(key, dsk[key])
+
+    assert set(result.keys()) == set(dsk.keys())
+    return result
+
+
+__all__ = ["compute_with_trace", "persist_with_trace"]

logxpy/decorators.py

@@ -0,0 +1,268 @@
+"""Logging decorators - @logged, @timed, @retry, @generator."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
+from functools import wraps
+from typing import Any, ParamSpec, TypeVar, cast
+
+from ._async import aaction, action
+from ._base import truncate
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+MASK_KEYS = {"password", "token", "secret", "key", "auth", "credential"}
+
+
+def _extract_args(
+    func: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+    capture_self: bool,
+    exclude: set[str],
+) -> dict[str, Any]:
+    sig = inspect.signature(func)
+    bound = sig.bind(*args, **kwargs)
+    bound.apply_defaults()
+    return {
+        k: "***" if k in exclude else truncate(v)
+        for k, v in bound.arguments.items()
+        if not (k == "self" and not capture_self)
+    }
+
+
+def logged(
+    fn: Callable[P, T] | None = None,
+    *,
+    level: str = "INFO",
+    capture_args: bool = True,
+    capture_result: bool = True,
+    capture_self: bool = False,
+    exclude: set[str] | None = None,
+    timer: bool = True,
+    when: Callable[..., bool] | None = None,
+    max_depth: int = 3,
+    max_length: int = 500,
+    silent_errors: bool = False,
+) -> Any:
+    """Universal logging decorator for entry/exit/timing/args/result."""
+    exclude = (exclude or set()) | MASK_KEYS
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        name = f"{func.__module__}.{func.__qualname__}"
+        is_async = inspect.iscoroutinefunction(func)
+
+        @wraps(func)
+        async def async_impl(*a: P.args, **kw: P.kwargs) -> T:
+            if when and not when(func, a, kw):
+                # We know func is async if we are here
+                return await cast("Awaitable[T]", func(*a, **kw))
+
+            log_args: dict[str, Any] = {}
+            if capture_args:
+                try:
+                    log_args = _extract_args(func, a, kw, capture_self, exclude)
+                except Exception:
+                    pass  # Best effort arg capture
+
+            async with aaction(name, level=level, **log_args) as act:
+                try:
+                    result = await cast("Awaitable[T]", func(*a, **kw))
+                    if capture_result and result is not None:
+                        act.fields["result"] = truncate(result, max_depth, max_length)
+                    return result
+                except Exception as e:
+                    if silent_errors:
+                        act.fields["error_suppressed"] = str(e)
+                        # We have to return something. Since we don't know T, we return None and cast it.
+                        return cast("T", None)
+                    raise
+
+        @wraps(func)
+        def sync_impl(*a: P.args, **kw: P.kwargs) -> T:
+            if when and not when(func, a, kw):
+                return func(*a, **kw)
+
+            log_args: dict[str, Any] = {}
+            if capture_args:
+                try:
+                    log_args = _extract_args(func, a, kw, capture_self, exclude)
+                except Exception:
+                    pass
+
+            with action(name, level=level, **log_args) as act:
+                try:
+                    result = func(*a, **kw)
+                    if capture_result and result is not None:
+                        act.fields["result"] = truncate(result, max_depth, max_length)
+                    return result
+                except Exception as e:
+                    if silent_errors:
+                        act.fields["error_suppressed"] = str(e)
+                        return cast("T", None)
+                    raise
+
+        return cast("Callable[P, T]", async_impl if is_async else sync_impl)
+
+    return decorator(fn) if fn else decorator
+
+
+def timed(metric: str | None = None) -> Any:
+    """Timing-only decorator."""
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        name = metric or f"{func.__module__}.{func.__qualname__}"
+        is_async = inspect.iscoroutinefunction(func)
+
+        @wraps(func)
+        async def async_impl(*a: P.args, **kw: P.kwargs) -> T:
+            start = time.monotonic()
+            try:
+                return await cast("Awaitable[T]", func(*a, **kw))
+            finally:
+                from .loggerx import log
+
+                log.debug(f"⏱ {func.__name__}", metric=name, duration_ms=round((time.monotonic() - start) * 1000, 2))
+
+        @wraps(func)
+        def sync_impl(*a: P.args, **kw: P.kwargs) -> T:
+            start = time.monotonic()
+            try:
+                return func(*a, **kw)
+            finally:
+                from .loggerx import log
+
+                log.debug(f"⏱ {func.__name__}", metric=name, duration_ms=round((time.monotonic() - start) * 1000, 2))
+
+        return cast("Callable[P, T]", async_impl if is_async else sync_impl)
+
+    return decorator
+
+
+def retry(
+    attempts: int = 3,
+    delay: float = 1.0,
+    backoff: float = 2.0,
+    on_retry: Callable[[int, Exception], None] | None = None,
+) -> Any:
+    """Retry with exponential backoff."""
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        is_async = inspect.iscoroutinefunction(func)
+
+        @wraps(func)
+        async def async_impl(*a: P.args, **kw: P.kwargs) -> T:
+            d, last = delay, None
+            for i in range(1, attempts + 1):
+                try:
+                    return await cast("Awaitable[T]", func(*a, **kw))
+                except Exception as e:
+                    last = e
+                    if i == attempts:
+                        raise
+                    if on_retry:
+                        on_retry(i, e)
+                    await asyncio.sleep(d)
+                    d *= backoff
+            # Should be unreachable if attempts > 0 and we raise on last attempt
+            raise last or Exception("Retry failed")
+
+        @wraps(func)
+        def sync_impl(*a: P.args, **kw: P.kwargs) -> T:
+            d, last = delay, None
+            for i in range(1, attempts + 1):
+                try:
+                    return func(*a, **kw)
+                except Exception as e:
+                    last = e
+                    if i == attempts:
+                        raise
+                    if on_retry:
+                        on_retry(i, e)
+                    time.sleep(d)
+                    d *= backoff
+            raise last or Exception("Retry failed")
+
+        return cast("Callable[P, T]", async_impl if is_async else sync_impl)
+
+    return decorator
+
+
+def generator(name: str | None = None, every: int = 100) -> Any:
+    """Generator progress tracking."""
+
+    def decorator(func: Callable[..., Iterator[T]]) -> Callable[..., Iterator[T]]:
+        label = name or func.__name__
+
+        @wraps(func)
+        def wrapper(*a: Any, **kw: Any) -> Iterator[T]:
+            from .loggerx import log
+
+            for i, item in enumerate(func(*a, **kw), 1):
+                if i % every == 0:
+                    log.info(f"📦 {label}", count=i)
+                yield item
+
+        return wrapper
+
+    return decorator
+
+
+def aiterator(name: str | None = None, every: int = 100) -> Any:
+    """Async iterator progress tracking."""
+
+    def decorator(func: Callable[..., AsyncIterator[T]]) -> Callable[..., AsyncIterator[T]]:
+        label = name or func.__name__
+
+        @wraps(func)
+        async def wrapper(*a: Any, **kw: Any) -> AsyncIterator[T]:
+            from .loggerx import log
+
+            i = 0
+            async for item in func(*a, **kw):
+                i += 1
+                if i % every == 0:
+                    log.info(f"📦 {label}", count=i)
+                yield item
+
+        return wrapper
+
+    return decorator
+
+
+def trace(name: str | None = None, kind: str = "internal", attributes: dict[str, Any] | None = None) -> Any:
+    """OpenTelemetry trace decorator."""
+
+    def decorator(func: Callable[P, T]) -> Callable[P, T]:
+        span_name = name or f"{func.__module__}.{func.__qualname__}"
+        is_async = inspect.iscoroutinefunction(func)
+
+        @wraps(func)
+        async def async_impl(*a: P.args, **kw: P.kwargs) -> T:
+            try:
+                from opentelemetry import trace as otel
+
+                tracer = otel.get_tracer(__name__)
+                with tracer.start_as_current_span(span_name, attributes=attributes):
+                    return await cast("Awaitable[T]", func(*a, **kw))
+            except ImportError:
+                return await cast("Awaitable[T]", func(*a, **kw))
+
+        @wraps(func)
+        def sync_impl(*a: P.args, **kw: P.kwargs) -> T:
+            try:
+                from opentelemetry import trace as otel
+
+                tracer = otel.get_tracer(__name__)
+                with tracer.start_as_current_span(span_name, attributes=attributes):
+                    return func(*a, **kw)
+            except ImportError:
+                return func(*a, **kw)
+
+        return cast("Callable[P, T]", async_impl if is_async else sync_impl)
+
+    return decorator

logxpy/filter.py

@@ -0,0 +1,124 @@
+"""
+Command line program for filtering line-based Eliot logs.
+"""
+
+if __name__ == "__main__":
+    import eliot.filter
+
+    eliot.filter.main()
+
+import sys
+from datetime import datetime, timedelta
+from json import JSONEncoder, dumps, loads
+
+
+class _DatetimeJSONEncoder(JSONEncoder):
+    """
+    JSON encoder that supports L{datetime}.
+    """
+
+    def default(self, o):
+        if isinstance(o, datetime):
+            return o.isoformat()
+        return JSONEncoder.default(self, o)
+
+
+class EliotFilter(object):
+    """
+    Filter Eliot log lines using a Python expression.
+
+    @ivar code: A Python code object, the compiled filter expression.
+    """
+
+    _SKIP = object()
+
+    def __init__(self, expr, incoming, output):
+        """
+        @param expr: A Python expression that will be called for each log message.
+        @type expr: L{str}
+
+        @param incoming: An iterable of L{bytes}, each of which is a serialized
+            Eliot message.
+
+        @param output: A file to which output should be written.
+        @type output: L{file} or a file-like object.
+        """
+        self.code = compile(expr, "<string>", "eval")
+        self.incoming = incoming
+        self.output = output
+
+    def run(self):
+        """
+        For each incoming message, decode the JSON, evaluate expression, encode
+        as JSON and write that to the output file.
+        """
+        for line in self.incoming:
+            message = loads(line)
+            result = self._evaluate(message)
+            if result is self._SKIP:
+                continue
+            self.output.write(dumps(result, cls=_DatetimeJSONEncoder) + "\n")
+
+    def _evaluate(self, message):
+        """
+        Evaluate the expression with the given Python object in its locals.
+
+        @param message: A decoded JSON input.
+
+        @return: The resulting object.
+        """
+        return eval(
+            self.code,
+            globals(),
+            {
+                "J": message,
+                "timedelta": timedelta,
+                "datetime": datetime,
+                "SKIP": self._SKIP,
+            },
+        )
+
+
+USAGE = """\
+Usage: cat eliot.log | python -m eliot.filter <expr>
+
+Read JSON-expression per line from stdin, and filter it using a Python
+expression <expr>.
+
+The expression will have a local `J` containing decoded JSON. `datetime` and
+`timedelta` from Python's `datetime` module are also available as locals,
+containing the corresponding classes. `SKIP` is also available, if it's the
+expression result that indicates nothing should be output.
+
+The output will be written to stdout using JSON serialization. `datetime`
+objects will be serialized to ISO format.
+
+Examples:
+
+- Pass through the messages unchanged:
+
+    $ cat eliot.log | python -m eliot.filter J
+
+- Retrieve a specific field from a specific message type, dropping messages
+  of other types:
+
+    $ cat eliot.log | python -m eliot.filter \\
+        "J['field'] if J.get('message_type') == 'my:message' else SKIP"
+"""
+
+
+def main(sys=sys):
+    """
+    Run the program.
+
+    Accept arguments from L{sys.argv}, read from L{sys.stdin}, write to
+    L{sys.stdout}.
+
+    @param sys: An object with same interface and defaulting to the L{sys}
+        module.
+    """
+    if len(sys.argv) != 2:
+        sys.stderr.write(USAGE)
+        return 1
+    EliotFilter(sys.argv[1], sys.stdin, sys.stdout).run()
+    return 0

logxpy/journald.py

@@ -0,0 +1,88 @@
+"""
+journald support for Eliot.
+"""
+
+from cffi import FFI
+from os import strerror
+from sys import argv
+from os.path import basename
+
+from .json import _dumps_bytes as dumps
+from ._message import TASK_UUID_FIELD, MESSAGE_TYPE_FIELD
+from ._action import ACTION_TYPE_FIELD, ACTION_STATUS_FIELD, FAILED_STATUS
+
+_ffi = FFI()
+_ffi.cdef(
+    """
+int sd_journal_send(const char *format, ...);
+"""
+)
+try:
+    try:
+        _journald = _ffi.dlopen("libsystemd.so.0")
+    except OSError:
+        # Older versions of systemd have separate library:
+        _journald = _ffi.dlopen("libsystemd-journal.so.0")
+except OSError as e:
+    raise ImportError("Failed to load journald: " + str(e))
+
+
+def sd_journal_send(**kwargs):
+    """
+    Send a message to the journald log.
+
+    @param kwargs: Mapping between field names to values, both as bytes.
+
+    @raise IOError: If the operation failed.
+    """
+    # The function uses printf formatting, so we need to quote
+    # percentages.
+    fields = [
+        _ffi.new("char[]", key.encode("ascii") + b"=" + value.replace(b"%", b"%%"))
+        for key, value in kwargs.items()
+    ]
+    fields.append(_ffi.NULL)
+    result = _journald.sd_journal_send(*fields)
+    if result != 0:
+        raise IOError(-result, strerror(-result))
+
+
+class JournaldDestination(object):
+    """
+    A logging destination that writes to journald.
+
+    The message will be logged as JSON, with an additional field
+    C{ELIOT_TASK} storing the C{task_uuid} and C{ELIOT_TYPE} storing the
+    C{message_type} or C{action_type}.
+
+    Messages for failed actions will get priority 3 ("error"), and
+    traceback messages will get priority 2 ("critical"). All other
+    messages will get priority 1 ("info").
+    """
+
+    def __init__(self):
+        self._identifier = basename(argv[0]).encode("utf-8")
+
+    def __call__(self, message):
+        """
+        Write the given message to journald.
+
+        @param message: Dictionary passed from a C{Logger}.
+        """
+        eliot_type = ""
+        priority = b"6"
+        if ACTION_TYPE_FIELD in message:
+            eliot_type = message[ACTION_TYPE_FIELD]
+            if message[ACTION_STATUS_FIELD] == FAILED_STATUS:
+                priority = b"3"
+        elif MESSAGE_TYPE_FIELD in message:
+            eliot_type = message[MESSAGE_TYPE_FIELD]
+            if eliot_type == "eliot:traceback":
+                priority = b"2"
+        sd_journal_send(
+            MESSAGE=dumps(message),
+            ELIOT_TASK=message[TASK_UUID_FIELD].encode("utf-8"),
+            ELIOT_TYPE=eliot_type.encode("utf-8"),
+            SYSLOG_IDENTIFIER=self._identifier,
+            PRIORITY=priority,
+        )