road24-artifacthub 0.1.0__py3-none-any.whl

road24_sdk/integrations/fastapi.py

@@ -0,0 +1,258 @@
+import json
+import logging
+import secrets
+import time
+from contextvars import ContextVar
+from http import HTTPStatus
+from typing import TYPE_CHECKING, Any
+
+from road24_sdk._formatter import clear_trace_context, get_trace_id, set_trace_context
+from road24_sdk._sanitizer import sanitize_body
+from road24_sdk._schemas import HttpAttributes
+from road24_sdk._types import HttpDirection, LogType
+from road24_sdk.integrations._base import Integration
+
+if TYPE_CHECKING:
+    from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+logger = logging.getLogger(__name__)
+
+_request_exception: ContextVar[Exception | None] = ContextVar("request_exception", default=None)
+
+
+class HttpInputLogger:
+    def log(
+        self,
+        scope: "Scope",
+        status_code: int,
+        duration_seconds: float,
+        request_body: bytes = b"",
+        response_body: bytes = b"",
+        record_metrics: bool = True,
+        exc: Exception | None = None,
+    ) -> None:
+        method = scope.get("method", "")
+        path = scope.get("path", "")
+        scheme = scope.get("scheme", "http")
+        headers = dict(scope.get("headers", []))
+        host = headers.get(b"host", b"").decode() or ""
+        url = f"{scheme}://{host}{path}" if host else path
+
+        if record_metrics:
+            from road24_sdk.metrics.http import record_http_request
+
+            record_http_request(
+                method=method,
+                url=url,
+                status_code=status_code,
+                duration_seconds=duration_seconds,
+                direction=HttpDirection.INPUT,
+            )
+
+        http_attrs = HttpAttributes(
+            direction=HttpDirection.INPUT,
+            method=method,
+            url=url,
+            target=path,
+            status_code=status_code,
+            duration_seconds=duration_seconds,
+            request_body=self._decode_body(request_body),
+            response_body=self._decode_body(response_body),
+            error_class=type(exc).__name__ if exc else "",
+            error_message=str(exc) if exc else "",
+        )
+
+        log_level = self._get_log_level(status_code)
+        getattr(logger, log_level)(
+            LogType.HTTP_REQUEST, exc_info=exc if exc else None, extra=http_attrs.as_dict()
+        )
+
+    def generate_trace_id(self, scope: "Scope") -> str:
+        headers = dict(scope.get("headers", []))
+        return headers.get(b"x-trace-id", b"").decode() or secrets.token_hex(16)
+
+    def _decode_body(self, body: bytes) -> str:
+        if not body:
+            return ""
+        try:
+            decoded = body.decode("utf-8")
+            return sanitize_body(decoded)
+        except UnicodeDecodeError:
+            return "<binary data>"
+
+    def _get_log_level(self, status_code: int) -> str:
+        if status_code < HTTPStatus.BAD_REQUEST:
+            return "info"
+        if status_code < HTTPStatus.INTERNAL_SERVER_ERROR:
+            return "warning"
+        return "error"
+
+
+class FastApiLoggingIntegration(Integration):
+    def __init__(
+        self,
+        *,
+        enable_logging: bool = True,
+        enable_metrics: bool = True,
+        enable_exception_handlers: bool = True,
+    ) -> None:
+        self._enable_logging = enable_logging
+        self._enable_metrics = enable_metrics
+        self._enable_exception_handlers = enable_exception_handlers
+
+    def setup_once(self) -> None:
+        pass
+
+    def setup(self, app: "ASGIApp") -> None:
+        if self._enable_exception_handlers and self._is_fastapi_app(app):
+            self._register_exception_handlers(app)
+
+        app.add_middleware(
+            _LogMiddleware,
+            enable_logging=self._enable_logging,
+            enable_metrics=self._enable_metrics,
+        )
+
+    @staticmethod
+    def _is_fastapi_app(app: Any) -> bool:
+        try:
+            from starlette.applications import Starlette
+        except ImportError:
+            return False
+        return isinstance(app, Starlette)
+
+    def _register_exception_handlers(self, app: Any) -> None:
+        from fastapi import HTTPException
+        from fastapi.exceptions import RequestValidationError
+
+        app.add_exception_handler(HTTPException, _http_exception_handler)
+        app.add_exception_handler(RequestValidationError, _validation_exception_handler)
+        app.add_exception_handler(Exception, _general_exception_handler)
+
+
+def _get_trace_id_for_response() -> str | None:
+    trace_id = get_trace_id()
+    return trace_id if trace_id != "0" * 32 else None
+
+
+async def _http_exception_handler(request: Any, exc: Any) -> Any:
+    from fastapi.responses import JSONResponse
+
+    if exc.status_code >= HTTPStatus.INTERNAL_SERVER_ERROR:
+        _request_exception.set(exc)
+
+    return JSONResponse(
+        status_code=exc.status_code,
+        content={"detail": exc.detail, "trace_id": _get_trace_id_for_response()},
+        headers=getattr(exc, "headers", None),
+    )
+
+
+async def _validation_exception_handler(request: Any, exc: Any) -> Any:
+    from fastapi.responses import JSONResponse
+
+    _request_exception.set(exc)
+
+    return JSONResponse(
+        status_code=HTTPStatus.UNPROCESSABLE_ENTITY.value,
+        content={"detail": exc.errors(), "trace_id": _get_trace_id_for_response()},
+    )
+
+
+async def _general_exception_handler(request: Any, exc: Exception) -> Any:
+    from fastapi.responses import JSONResponse
+
+    _request_exception.set(exc)
+
+    return JSONResponse(
+        status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+        content={
+            "detail": "Internal server error",
+            "trace_id": _get_trace_id_for_response(),
+            "error_type": type(exc).__name__,
+        },
+    )
+
+
+class _LogMiddleware:
+    def __init__(
+        self,
+        app: "ASGIApp",
+        *,
+        enable_logging: bool = True,
+        enable_metrics: bool = True,
+    ) -> None:
+        self.app = app
+        self._http_logger = HttpInputLogger()
+        self._enable_logging = enable_logging
+        self._enable_metrics = enable_metrics
+
+    async def __call__(
+        self,
+        scope: "Scope",
+        receive: "Receive",
+        send: "Send",
+    ) -> None:
+        if scope["type"] != "http" or scope.get("path") == "/metrics":
+            await self.app(scope, receive, send)
+            return
+
+        trace_id = self._http_logger.generate_trace_id(scope)
+        trace_token = set_trace_context(trace_id)
+
+        request_body = b""
+        response_body = b""
+        status_code = 500
+        start_time = time.perf_counter()
+
+        async def receive_wrapper() -> "Message":
+            nonlocal request_body
+            message = await receive()
+            if message["type"] == "http.request":
+                request_body += message.get("body", b"")
+            return message
+
+        async def send_wrapper(message: "Message") -> None:
+            nonlocal response_body, status_code
+            if message["type"] == "http.response.start":
+                status_code = message["status"]
+                headers = [
+                    *message.get("headers", []),
+                    (b"x-trace-id", trace_id.encode()),
+                ]
+                message = {**message, "headers": headers}
+            elif message["type"] == "http.response.body":
+                response_body += message.get("body", b"")
+            await send(message)
+
+        try:
+            await self.app(scope, receive_wrapper, send_wrapper)
+        except Exception as e:
+            _request_exception.set(e)
+            status_code = HTTPStatus.INTERNAL_SERVER_ERROR
+            error_body = json.dumps({
+                "detail": "Internal server error",
+                "trace_id": trace_id,
+                "error_type": type(e).__name__,
+            }).encode()
+            await send({"type": "http.response.start", "status": status_code, "headers": [
+                (b"content-type", b"application/json"),
+                (b"x-trace-id", trace_id.encode()),
+            ]})
+            await send({"type": "http.response.body", "body": error_body})
+            response_body = error_body
+        finally:
+            exc = _request_exception.get()
+            _request_exception.set(None)
+            duration_seconds = time.perf_counter() - start_time
+            if self._enable_logging:
+                self._http_logger.log(
+                    scope=scope,
+                    status_code=status_code,
+                    duration_seconds=duration_seconds,
+                    request_body=request_body,
+                    response_body=response_body,
+                    record_metrics=self._enable_metrics,
+                    exc=exc,
+                )
+            clear_trace_context(trace_token)

road24_sdk/integrations/httpx.py

@@ -0,0 +1,120 @@
+import logging
+import time
+from http import HTTPStatus
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
+
+from road24_sdk._sanitizer import sanitize_body
+from road24_sdk._schemas import HttpAttributes
+from road24_sdk._types import HttpDirection, LogType
+from road24_sdk.integrations._base import Integration
+
+if TYPE_CHECKING:
+    from httpx import AsyncClient, Request, Response
+
+logger = logging.getLogger(__name__)
+
+
+class HttpOutputLogger:
+    async def log_request_hook(self, request: "Request") -> None:
+        request.extensions["start_time"] = time.perf_counter()
+
+    async def log_response_hook(self, response: "Response") -> None:
+        await response.aread()
+        start_time = response.request.extensions.get("start_time")
+        duration_seconds = time.perf_counter() - start_time if start_time else 0.0
+        self.log(response.request, response, duration_seconds)
+
+    def log(
+        self,
+        request: "Request",
+        response: "Response",
+        duration_seconds: float,
+        record_metrics: bool = True,
+    ) -> None:
+        parsed = urlparse(str(request.url))
+        url = str(request.url)
+
+        if record_metrics:
+            from road24_sdk.metrics.http import record_http_request
+
+            record_http_request(
+                method=request.method,
+                url=url,
+                status_code=response.status_code,
+                duration_seconds=duration_seconds,
+                direction=HttpDirection.OUTPUT,
+            )
+
+        request_body = request.content.decode("utf-8") if request.content else ""
+        response_body = response.text if response.text else ""
+
+        http_attrs = HttpAttributes(
+            direction=HttpDirection.OUTPUT,
+            method=request.method,
+            url=url,
+            target=parsed.path,
+            status_code=response.status_code,
+            duration_seconds=duration_seconds,
+            request_body=sanitize_body(request_body),
+            response_body=sanitize_body(response_body),
+        )
+
+        log_level = self._get_log_level(response.status_code)
+        getattr(logger, log_level)(LogType.HTTP_REQUEST, extra=http_attrs.as_dict())
+
+    def _get_log_level(self, status_code: int) -> str:
+        if status_code < HTTPStatus.BAD_REQUEST:
+            return "info"
+        if status_code < HTTPStatus.INTERNAL_SERVER_ERROR:
+            return "warning"
+        return "error"
+
+
+class HttpxLoggingIntegration(Integration):
+    def __init__(
+        self,
+        *,
+        enable_logging: bool = True,
+        enable_metrics: bool = True,
+    ) -> None:
+        self._enable_logging = enable_logging
+        self._enable_metrics = enable_metrics
+        self._output_logger = HttpOutputLogger()
+
+    def setup_once(self) -> None:
+        if self._mark_setup_once():
+            return
+        from httpx import AsyncClient
+
+        original_init = AsyncClient.__init__
+        output_logger = self._output_logger
+        enable_logging = self._enable_logging
+
+        def _patched_init(self_client: "AsyncClient", *args: Any, **kwargs: Any) -> None:
+            original_init(self_client, *args, **kwargs)
+            if enable_logging:
+                hooks = self_client.event_hooks
+                hooks["request"].append(output_logger.log_request_hook)
+                hooks["response"].append(output_logger.log_response_hook)
+
+        AsyncClient.__init__ = _patched_init  # type: ignore[method-assign]
+
+    def setup(self, client: "AsyncClient") -> None:
+        if self._enable_logging:
+            hooks = client.event_hooks
+            hooks["request"].append(self._output_logger.log_request_hook)
+            hooks["response"].append(self._output_logger.log_response_hook)
+
+    def create_client(self, **kwargs: Any) -> "AsyncClient":
+        from httpx import AsyncClient
+
+        event_hooks: dict[str, list[Any]] = kwargs.pop(
+            "event_hooks", {"request": [], "response": []}
+        )
+
+        if self._enable_logging:
+            event_hooks["request"].append(self._output_logger.log_request_hook)
+            event_hooks["response"].append(self._output_logger.log_response_hook)
+
+        return AsyncClient(event_hooks=event_hooks, **kwargs)

road24_sdk/integrations/redis.py

@@ -0,0 +1,233 @@
+import logging
+import time
+from typing import TYPE_CHECKING, Any
+
+from road24_sdk._schemas import RedisAttributes
+from road24_sdk._types import LogType, RedisOperation
+from road24_sdk.integrations._base import Integration
+
+if TYPE_CHECKING:
+    from redis.asyncio import Redis
+
+logger = logging.getLogger(__name__)
+
+OPERATION_MAP: dict[str, RedisOperation] = {
+    "get": RedisOperation.GET,
+    "set": RedisOperation.SET,
+    "delete": RedisOperation.DELETE,
+    "del": RedisOperation.DELETE,
+    "expire": RedisOperation.EXPIRE,
+    "hget": RedisOperation.HGET,
+    "hset": RedisOperation.HSET,
+    "lpush": RedisOperation.LPUSH,
+    "rpush": RedisOperation.RPUSH,
+    "lpop": RedisOperation.LPOP,
+    "rpop": RedisOperation.RPOP,
+    "publish": RedisOperation.PUBLISH,
+    "subscribe": RedisOperation.SUBSCRIBE,
+}
+
+
+class RedisLogger:
+    def __init__(self, *, record_metrics: bool = True) -> None:
+        self._record_metrics = record_metrics
+
+    def log(
+        self,
+        operation: RedisOperation,
+        key: str,
+        duration_seconds: float,
+        args: str = "",
+    ) -> None:
+        if self._record_metrics:
+            from road24_sdk.metrics.redis import record_redis_command
+
+            record_redis_command(
+                operation=operation,
+                key=key,
+                duration_seconds=duration_seconds,
+            )
+
+        redis_attrs = RedisAttributes(
+            operation=operation,
+            key=key,
+            duration_seconds=duration_seconds,
+            command_args=args,
+        )
+
+        logger.info(LogType.REDIS_COMMAND, extra=redis_attrs.as_dict())
+
+
+class LoggedRedis:
+    def __init__(
+        self,
+        client: "Redis",
+        *,
+        record_metrics: bool = True,
+    ) -> None:
+        self._client = client
+        self._logger = RedisLogger(record_metrics=record_metrics)
+
+    async def get(self, key: str) -> Any:
+        return await self._execute("get", key)
+
+    async def set(
+        self,
+        key: str,
+        value: Any,
+        ex: int | None = None,
+        px: int | None = None,
+        nx: bool = False,
+        xx: bool = False,
+    ) -> Any:
+        kwargs: dict[str, Any] = {}
+        if ex is not None:
+            kwargs["ex"] = ex
+        if px is not None:
+            kwargs["px"] = px
+        if nx:
+            kwargs["nx"] = nx
+        if xx:
+            kwargs["xx"] = xx
+
+        args_str = f"value={value}"
+        if kwargs:
+            args_str += ", " + ", ".join(f"{k}={v}" for k, v in kwargs.items())
+
+        return await self._execute("set", key, value, args_str=args_str, **kwargs)
+
+    async def delete(self, *keys: str) -> int:
+        key = keys[0] if keys else "unknown"
+        return await self._execute("delete", key, *keys)
+
+    async def expire(self, key: str, seconds: int) -> bool:
+        return await self._execute("expire", key, seconds, args_str=f"seconds={seconds}")
+
+    async def hget(self, name: str, key: str) -> Any:
+        return await self._execute("hget", name, key, args_str=f"field={key}")
+
+    async def hset(
+        self,
+        name: str,
+        key: str | None = None,
+        value: Any = None,
+        mapping: dict[str, Any] | None = None,
+    ) -> int:
+        args_str = f"field={key}" if key else f"mapping={len(mapping or {})} fields"
+        return await self._execute(
+            "hset",
+            name,
+            key,
+            value,
+            mapping=mapping,
+            args_str=args_str,
+        )
+
+    async def lpush(self, key: str, *values: Any) -> int:
+        return await self._execute("lpush", key, *values, args_str=f"count={len(values)}")
+
+    async def rpush(self, key: str, *values: Any) -> int:
+        return await self._execute("rpush", key, *values, args_str=f"count={len(values)}")
+
+    async def lpop(self, key: str, count: int | None = None) -> Any:
+        return await self._execute("lpop", key, count=count)
+
+    async def rpop(self, key: str, count: int | None = None) -> Any:
+        return await self._execute("rpop", key, count=count)
+
+    async def ping(self) -> bool:
+        return await self._client.ping()
+
+    async def aclose(self) -> None:
+        await self._client.aclose()
+
+    async def _execute(
+        self,
+        command: str,
+        key: str,
+        *args: Any,
+        args_str: str = "",
+        **kwargs: Any,
+    ) -> Any:
+        operation = OPERATION_MAP.get(command, RedisOperation.OTHER)
+        start_time = time.perf_counter()
+        try:
+            method = getattr(self._client, command)
+            return await method(key, *args, **kwargs)
+        finally:
+            duration = time.perf_counter() - start_time
+            self._logger.log(
+                operation=operation,
+                key=key,
+                duration_seconds=duration,
+                args=args_str,
+            )
+
+
+class RedisLoggingIntegration(Integration):
+    def __init__(
+        self,
+        *,
+        enable_logging: bool = True,
+        enable_metrics: bool = True,
+    ) -> None:
+        self._enable_logging = enable_logging
+        self._enable_metrics = enable_metrics
+
+    def setup_once(self) -> None:
+        if self._mark_setup_once():
+            return
+        from redis.asyncio import Redis
+
+        original_execute = Redis.execute_command
+        redis_logger = RedisLogger(record_metrics=self._enable_metrics)
+        enable_logging = self._enable_logging
+
+        async def _patched_execute(
+            self_client: "Redis", *args: Any, **kwargs: Any,
+        ) -> Any:
+            command = str(args[0]).lower() if args else "unknown"
+            key = str(args[1]) if len(args) > 1 else ""
+            operation = OPERATION_MAP.get(command, RedisOperation.OTHER)
+            start_time = time.perf_counter()
+            try:
+                return await original_execute(self_client, *args, **kwargs)
+            finally:
+                duration = time.perf_counter() - start_time
+                if enable_logging:
+                    redis_logger.log(
+                        operation=operation,
+                        key=key,
+                        duration_seconds=duration,
+                    )
+
+        Redis.execute_command = _patched_execute  # type: ignore[method-assign]
+
+    def setup(self, client: "Redis") -> None:
+        original_execute = client.execute_command
+        redis_logger = RedisLogger(record_metrics=self._enable_metrics)
+        enable_logging = self._enable_logging
+
+        async def _patched_execute(*args: Any, **kwargs: Any) -> Any:
+            command = str(args[0]).lower() if args else "unknown"
+            key = str(args[1]) if len(args) > 1 else ""
+            operation = OPERATION_MAP.get(command, RedisOperation.OTHER)
+            start_time = time.perf_counter()
+            try:
+                return await original_execute(*args, **kwargs)
+            finally:
+                duration = time.perf_counter() - start_time
+                if enable_logging:
+                    redis_logger.log(
+                        operation=operation,
+                        key=key,
+                        duration_seconds=duration,
+                    )
+
+        client.execute_command = _patched_execute  # type: ignore[method-assign]
+
+    def instrument(self, client: "Redis") -> LoggedRedis:
+        return LoggedRedis(
+            client,
+            record_metrics=self._enable_metrics,
+        )