road24-artifacthub 0.1.0__py3-none-any.whl

road24_sdk/integrations/sqlalchemy.py

@@ -0,0 +1,175 @@
+import logging
+import re
+import time
+from typing import Any
+
+from road24_sdk._schemas import DbAttributes
+from road24_sdk._types import DbOperation, LogType
+from road24_sdk.integrations._base import Integration
+
+logger = logging.getLogger(__name__)
+
+_START_TIME_KEY = "_query_start_time"
+
+
+class DbLogger:
+    MAX_STATEMENT_LENGTH = 500
+
+    TABLE_PATTERNS = [
+        r"FROM\s+([\"']?[\w\.]+[\"']?)",
+        r"INTO\s+([\"']?[\w\.]+[\"']?)",
+        r"UPDATE\s+([\"']?[\w\.]+[\"']?)",
+        r"DELETE\s+FROM\s+([\"']?[\w\.]+[\"']?)",
+    ]
+
+    def log_query(
+        self,
+        operation: DbOperation,
+        table: str,
+        duration_seconds: float,
+        statement: str = "",
+        record_metrics: bool = True,
+    ) -> None:
+        if record_metrics:
+            from road24_sdk.metrics.db import record_db_query
+
+            record_db_query(
+                operation=operation,
+                table=table,
+                duration_seconds=duration_seconds,
+            )
+
+        db_attrs = DbAttributes(
+            operation=operation,
+            table=table,
+            duration_seconds=duration_seconds,
+            statement=statement[: self.MAX_STATEMENT_LENGTH],
+        )
+
+        logger.info(LogType.DB_QUERY, extra=db_attrs.as_dict())
+
+    def before_execute(
+        self,
+        conn: Any,
+        cursor: Any,
+        statement: str,
+        parameters: Any,
+        context: Any,
+        executemany: bool,
+    ) -> None:
+        conn.info[_START_TIME_KEY] = time.perf_counter()
+
+    def after_execute(
+        self,
+        conn: Any,
+        cursor: Any,
+        statement: str,
+        parameters: Any,
+        context: Any,
+        executemany: bool,
+        record_metrics: bool = True,
+    ) -> None:
+        start_time = conn.info.pop(_START_TIME_KEY, None)
+        duration_seconds = time.perf_counter() - start_time if start_time else 0.0
+        operation = self.extract_operation(statement)
+        table = self.extract_table(statement)
+
+        self.log_query(
+            operation=operation,
+            table=table,
+            duration_seconds=duration_seconds,
+            statement=statement,
+            record_metrics=record_metrics,
+        )
+
+    def extract_operation(self, statement: str) -> DbOperation:
+        statement_upper = statement.strip().upper()
+
+        if statement_upper.startswith("SELECT"):
+            return DbOperation.SELECT
+        if statement_upper.startswith("INSERT"):
+            return DbOperation.INSERT
+        if statement_upper.startswith("UPDATE"):
+            return DbOperation.UPDATE
+        if statement_upper.startswith("DELETE"):
+            return DbOperation.DELETE
+
+        return DbOperation.OTHER
+
+    def extract_table(self, statement: str) -> str:
+        statement_upper = statement.strip().upper()
+
+        for pattern in self.TABLE_PATTERNS:
+            if match := re.search(pattern, statement_upper):
+                table = match.group(1).strip("\"'").lower()
+                return table.split(".")[-1]
+
+        return "unknown"
+
+
+class SqlalchemyLoggingIntegration(Integration):
+    def __init__(
+        self,
+        *,
+        enable_logging: bool = True,
+        enable_metrics: bool = True,
+    ) -> None:
+        self._enable_logging = enable_logging
+        self._enable_metrics = enable_metrics
+        self._db_logger = DbLogger()
+
+    def setup_once(self) -> None:
+        if self._mark_setup_once():
+            return
+        from sqlalchemy import event
+        from sqlalchemy.engine import Engine
+
+        event.listen(Engine, "before_cursor_execute", self._db_logger.before_execute)
+
+        if self._enable_logging:
+            enable_metrics = self._enable_metrics
+            db_logger = self._db_logger
+
+            event.listen(
+                Engine,
+                "after_cursor_execute",
+                lambda conn, cursor, stmt, params, ctx, many: (
+                    db_logger.after_execute(
+                        conn,
+                        cursor,
+                        stmt,
+                        params,
+                        ctx,
+                        many,
+                        record_metrics=enable_metrics,
+                    )
+                ),
+            )
+
+    def setup(self, engine: Any) -> None:
+        from sqlalchemy import event
+
+        sync_engine = getattr(engine, "sync_engine", engine)
+
+        event.listen(
+            sync_engine,
+            "before_cursor_execute",
+            self._db_logger.before_execute,
+        )
+
+        if self._enable_logging:
+            event.listen(
+                sync_engine,
+                "after_cursor_execute",
+                lambda conn, cursor, stmt, params, ctx, many: (
+                    self._db_logger.after_execute(
+                        conn,
+                        cursor,
+                        stmt,
+                        params,
+                        ctx,
+                        many,
+                        record_metrics=self._enable_metrics,
+                    )
+                ),
+            )

road24_sdk/metrics/__init__.py

@@ -0,0 +1,9 @@
+from road24_sdk.metrics.db import record_db_query
+from road24_sdk.metrics.http import record_http_request
+from road24_sdk.metrics.redis import record_redis_command
+
+__all__ = [
+    "record_db_query",
+    "record_http_request",
+    "record_redis_command",
+]

road24_sdk/metrics/db.py

@@ -0,0 +1,29 @@
+from prometheus_client import Counter, Histogram
+
+from road24_sdk._types import DbOperation
+
+DB_QUERY_DURATION = Histogram(
+    name="db_query_duration_seconds",
+    documentation="Database query duration in seconds",
+    labelnames=["operation", "table"],
+    buckets=(0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5),
+)
+
+DB_QUERY_TOTAL = Counter(
+    name="db_queries_total",
+    documentation="Total number of database queries",
+    labelnames=["operation", "table"],
+)
+
+
+def record_db_query(
+    operation: DbOperation,
+    table: str,
+    duration_seconds: float,
+) -> None:
+    labels = {
+        "operation": operation.value,
+        "table": table,
+    }
+    DB_QUERY_TOTAL.labels(**labels).inc()
+    DB_QUERY_DURATION.labels(**labels).observe(duration_seconds)

road24_sdk/metrics/http.py

@@ -0,0 +1,33 @@
+from prometheus_client import Counter, Histogram
+
+from road24_sdk._types import HttpDirection
+
+HTTP_REQUEST_DURATION = Histogram(
+    name="http_request_duration_seconds",
+    documentation="HTTP request duration in seconds",
+    labelnames=["method", "url", "status_code", "direction"],
+    buckets=(0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0),
+)
+
+HTTP_REQUEST_TOTAL = Counter(
+    name="http_requests_total",
+    documentation="Total number of HTTP requests",
+    labelnames=["method", "url", "status_code", "direction"],
+)
+
+
+def record_http_request(
+    method: str,
+    url: str,
+    status_code: int,
+    duration_seconds: float,
+    direction: HttpDirection,
+) -> None:
+    labels = {
+        "method": method,
+        "url": url,
+        "status_code": str(status_code),
+        "direction": direction.value,
+    }
+    HTTP_REQUEST_TOTAL.labels(**labels).inc()
+    HTTP_REQUEST_DURATION.labels(**labels).observe(duration_seconds)

road24_sdk/metrics/redis.py

@@ -0,0 +1,29 @@
+from prometheus_client import Counter, Histogram
+
+from road24_sdk._types import RedisOperation
+
+REDIS_COMMAND_DURATION = Histogram(
+    name="redis_command_duration_seconds",
+    documentation="Redis command duration in seconds",
+    labelnames=["operation", "key"],
+    buckets=(0.0005, 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5),
+)
+
+REDIS_COMMAND_TOTAL = Counter(
+    name="redis_commands_total",
+    documentation="Total number of Redis commands",
+    labelnames=["operation", "key"],
+)
+
+
+def record_redis_command(
+    operation: RedisOperation,
+    key: str,
+    duration_seconds: float,
+) -> None:
+    labels = {
+        "operation": operation.value,
+        "key": key,
+    }
+    REDIS_COMMAND_TOTAL.labels(**labels).inc()
+    REDIS_COMMAND_DURATION.labels(**labels).observe(duration_seconds)