writer 0.8.3rc4__py3-none-any.whl → 1.25.1rc1__py3-none-any.whl

writer/logs.py

@@ -0,0 +1,277 @@
+import io
+import json
+import logging
+import logging.config
+import os
+from contextlib import contextmanager, redirect_stdout
+from time import time
+from typing import Any, Callable, Dict, List, Optional, Union
+
+WRITER_LOG_LEVEL = os.getenv("WRITER_LOG_LEVEL", "INFO")
+WRITER_LOG_FORMAT = os.getenv("WRITER_LOG_FORMAT", "text")  # 'text' or 'json'
+
+FAILOVER_ROUTING_KEY = "unset"
+FAILOVER_BUFFER = "failover"
+
+
+def get_routing_key(prefix: Optional[str] = None) -> str:
+    try:
+        from writer.blueprints import get_current_block
+        current_block = get_current_block()
+    except RuntimeError:
+        current_block = None
+
+    key = FAILOVER_ROUTING_KEY
+    if current_block is not None:
+        key = current_block.component.id
+    return f"{prefix}-{key}"
+
+
+class RoutingMap():
+    """
+    Maintains a map of routing keys to in-memory output buffers (io.StringIO).
+    Used for capturing logs or stdout output in different contexts.
+    """
+
+    def __init__(self) -> None:
+        # It's not expected that this will be used without context.
+        # But just in case a fail-over buffer is provided
+        self._buffer_map: Dict[str, io.StringIO] = {
+            FAILOVER_BUFFER: io.StringIO(),
+        }
+    
+    def get_buffer(self, key: str) -> io.StringIO:
+        """
+        Retrieve the buffer associated with a routing key.
+
+        If buffer for the key is not present uses a fail-over buffer
+        """
+        return self._buffer_map.get(key, self._buffer_map[FAILOVER_BUFFER])
+
+    def add_buffer(self, key: str) -> io.StringIO:
+        """
+        Add a new buffer with the given key.
+        """
+        buffer = io.StringIO()
+        self._buffer_map[key] = buffer
+        return buffer
+
+    def remove_buffer(self, key: str) -> None:
+        """
+        Remove the buffer associated with the specified key.
+        """
+        self._buffer_map.pop(key, None)
+
+
+routing_map = RoutingMap()
+
+
+class RoutingStream(io.StringIO):
+    """
+    Custom stream that re-routes stdout to the correct io.StringIO buffer
+    based on the current context routing key.
+    """
+
+    def write(self, s: str) -> int:
+        if s.strip():
+            logging.getLogger("stdout").info(s)
+        routing_key = get_routing_key(prefix="stdout")
+        return routing_map.get_buffer(routing_key).write(s)
+    
+    def getvalue(self) -> str:
+        routing_key = get_routing_key(prefix="stdout")
+        return routing_map.get_buffer(routing_key).getvalue()
+
+
+class RoutingHandler(logging.StreamHandler):
+    """
+    Custom logging handler that re-routes logs to different buffers
+    based on the current context routing key.
+
+    Overwritten methods are mirroring original ones from logging.StreamHandler.
+    The only difference is how 'stream' object is acquired
+    """
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            routing_key = get_routing_key(prefix="logging")
+            stream = routing_map.get_buffer(routing_key)
+            stream.write(msg + self.terminator)
+            self.flush()
+        except RecursionError:
+            raise
+        except Exception:
+            self.handleError(record)
+
+    def flush(self):
+        routing_key = get_routing_key(prefix="logging")
+        stream = routing_map.get_buffer(routing_key)
+        with self.lock:
+            if stream and hasattr(stream, "flush"):
+                stream.flush()
+
+
+@contextmanager
+def use_stdout_redirect(add_log_entry_func: Union[Callable[[str], None], List[Callable[[str], None]]]):
+    """
+    Context manager that redirects stdout to a context-specific buffer.
+    Supports single callback or list of callbacks for multiple output destinations.
+    """
+    # Normalize to list of callbacks
+    callbacks = [add_log_entry_func] if callable(add_log_entry_func) else add_log_entry_func
+
+    key = get_routing_key(prefix="stdout")
+    buffer = routing_map.add_buffer(key)
+    try:
+        with redirect_stdout(RoutingStream()):
+            yield buffer
+    finally:
+        routing_map.remove_buffer(key)
+        stdout = buffer.getvalue()
+        if stdout:
+            for callback in callbacks:
+                callback(stdout)
+
+
+@contextmanager
+def use_logging_redirect(add_log_entry_func: Union[Callable[[str], None], List[Callable[[str], None]]]):
+    """
+    Context manager that redirects logging to a context-specific buffer.
+    Supports single callback or list of callbacks for multiple output destinations.
+    """
+    # Normalize to list of callbacks
+    callbacks = [add_log_entry_func] if callable(add_log_entry_func) else add_log_entry_func
+
+    key = get_routing_key(prefix="logging")
+    buffer = routing_map.add_buffer(key)
+    try:
+        yield buffer
+    finally:
+        routing_map.remove_buffer(key)
+        logs = buffer.getvalue()
+        if logs:
+            for callback in callbacks:
+                callback(logs)
+
+
+class JSONFormatter(logging.Formatter):
+    def format(self, record: logging.LogRecord, **kwargs) -> str:
+        try:
+            from writer.blueprints import get_current_block
+            from writer.core import get_app_process, get_session
+            current_block = get_current_block()
+            app_process = get_app_process()
+            session = get_session()
+        except RuntimeError:
+            current_block = None
+            session = None
+            app_process = None
+
+        data: Dict[str, Any] = {
+            "severity": record.levelname.upper(),
+            "message": super().format(record),
+            "logger_name": record.name,
+            "process": {
+                "name": record.processName
+            },
+            "timestamp": time(),
+        }
+
+        if current_block is not None:
+            data["component"] = {
+                "id": current_block.component.id,
+                "type": current_block.component.type
+            }
+
+        if session is not None:
+            data["session"] = {
+                "id": session.session_id,
+            }
+
+        if app_process is not None:
+            data["process"]["mode"] = app_process.mode
+
+        if isinstance(record.args, dict):
+            data.update(record.args)
+
+        # if record.args[0] is a dict add to the json dict
+        if isinstance(record.args, tuple):
+            if len(record.args) > 0 and isinstance(record.args[0], dict):
+                data.update(record.args[0])
+
+        try:
+            data_as_json = json.dumps(data)
+            return data_as_json 
+        except Exception:
+            return '{"unserializable": true}' 
+
+
+def get_handler(format: str = WRITER_LOG_FORMAT, level: str = WRITER_LOG_LEVEL):
+    return {
+        "level": level,
+        "class": "logging.StreamHandler",
+        "formatter": format,
+    }
+
+
+def get_logger(level: str = WRITER_LOG_LEVEL):
+    return {
+        "handlers": ["basic"],
+        "level": level,
+        "propagate": False,
+    }
+
+
+LOGGING_CONFIG: Dict[str, Any] = {
+    "version": 1,
+    "disable_existing_loggers": False,
+    "formatters": {
+        "text": {
+            "format": "%(levelname)s - %(name)s - %(message)s",
+        },
+        "json": {
+            "()": JSONFormatter,
+        },
+        "user": {
+            "format": "%(levelname)s - %(message)s"
+        }
+    },
+    "handlers": {
+        "basic": get_handler(),
+        "stdout": {
+            "level": "DEBUG",
+            "class": "logging.StreamHandler",
+            "formatter": "json" if WRITER_LOG_FORMAT == "json" else None
+        },
+        "routing": {
+            "()": RoutingHandler,
+            "formatter": "user",
+        }
+    },
+    "loggers": {
+        "root": get_logger(),
+        "writer": get_logger(),
+        "app": get_logger(),
+        "from_app": get_logger(),
+        "journal": get_logger(),
+        "kv_storage": get_logger(),
+        "vault": get_logger(),
+        "exec_logger": {
+            "handlers": ["basic", "routing"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
+        "user_code": {
+            "handlers": ["basic", "routing"],
+            "level": "DEBUG",
+            "propagate": False,
+        },
+        "stdout": {
+            "handlers": ["stdout"],
+            "propagate": False,
+        }
+    }
+}
+
+logging.config.dictConfig(LOGGING_CONFIG)