pyquerytracker 0.1.0__py3-none-any.whl → 0.1.1__py3-none-any.whl

pyquerytracker/core.py

@@ -1,86 +1,156 @@
+import asyncio
 import time
-import logging
 from functools import update_wrapper
-from typing import Any, Callable, TypeVar, Generic
+from typing import Any, Callable, Generic, Optional, TypeVar
+
 from pyquerytracker.config import get_config
+from pyquerytracker.db.writer import DBWriter
+from pyquerytracker.exporter.base import NullExporter
+from pyquerytracker.exporter.manager import ExporterManager
+from pyquerytracker.tracker import store_tracked_query
+from pyquerytracker.utils.logger import QueryLogger
 
-# Set up logger
-logger = logging.getLogger("pyquerytracker")
-if not logger.handlers:
-    handler = logging.StreamHandler()
-    formatter = logging.Formatter(
-        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
-    )
-    handler.setFormatter(formatter)
-    logger.addHandler(handler)
-    logger.setLevel(logging.INFO)
+logger = QueryLogger.get_logger()
 
 T = TypeVar("T")
 
 
 class TrackQuery(Generic[T]):
-    """
-    Class-based decorator to track and log the execution time of functions or methods.
-
-    Logs include:
-    - Function name
-    - Class name (if method)
-    - Execution time (ms)
-    - Arguments
-    - Errors (if any)
-
-    Usage:
-        @TrackQuery()
-        def my_function():
-            ...
-    """
-
     def __init__(self) -> None:
         self.config = get_config()
+        if self.config.export_type and self.config.export_path:
+            exporter = ExporterManager.create_exporter(self.config)
+            ExporterManager.set(exporter)
+            self.exporter = exporter
+        else:
+            self.exporter = NullExporter()
+
+    def _extract_class_name(self, args: Any) -> Optional[str]:
+        if args:
+            obj = args[0]
+            if hasattr(obj, "__class__"):
+                return obj.__name__ if isinstance(obj, type) else obj.__class__.__name__
+        return None
+
+    # pylint: disable=too-many-positional-arguments
+    def _build_log_data(self, func, class_name, duration, args, kwargs, error=None):
+        data = {
+            "event": (
+                "error"
+                if error
+                else (
+                    "slow_execution"
+                    if duration > self.config.slow_log_threshold_ms
+                    else "normal_execution"
+                )
+            ),
+            "function_name": func.__name__,
+            "class_name": class_name,
+            "duration_ms": duration,
+            "func_args": repr(args),
+            "func_kwargs": repr(kwargs),
+        }
+        if error:
+            data["error"] = str(error)
+        return data
+
+    def _handle_export(self, log_data):
+        self.exporter.append(log_data)
+        if self.config.persist_to_db:
+            DBWriter.save(log_data)
+        store_tracked_query(log_data)
 
     def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
+        if asyncio.iscoroutinefunction(func):
+
+            async def async_wrapped(*args: Any, **kwargs: Any) -> T:
+                start = time.perf_counter()
+                class_name = self._extract_class_name(args)
+
+                try:
+                    result = await func(*args, **kwargs)
+                    duration = (time.perf_counter() - start) * 1000
+                    log_data = self._build_log_data(
+                        func, class_name, duration, args, kwargs
+                    )
+
+                    if duration > self.config.slow_log_threshold_ms:
+                        logger.log(
+                            self.config.slow_log_level,
+                            "%s%s -> Slow execution: took %.2fms",
+                            f"{class_name}." if class_name else "",
+                            func.__name__,
+                            duration,
+                            extra=log_data,
+                        )
+                    else:
+                        logger.info(
+                            "Function %s%s executed successfully in %.2fms",
+                            f"{class_name}." if class_name else "",
+                            func.__name__,
+                            duration,
+                            extra=log_data,
+                        )
+
+                    self._handle_export(log_data)
+                    return result
+
+                except Exception as e:
+                    duration = (time.perf_counter() - start) * 1000
+                    log_data = self._build_log_data(
+                        func, class_name, duration, args, kwargs, error=e
+                    )
+                    logger.error(
+                        "Function %s%s failed after %.2fms: %s",
+                        f"{class_name}." if class_name else "",
+                        func.__name__,
+                        duration,
+                        str(e),
+                        exc_info=True,
+                        extra=log_data,
+                    )
+                    self._handle_export(log_data)
+                    return None
+
+            return update_wrapper(async_wrapped, func)
+
         def wrapped(*args: Any, **kwargs: Any) -> T:
             start = time.perf_counter()
-            class_name = None
-
-            # Try to detect if this is an instance or class method
-            if args:
-                possible_self_or_cls = args[0]
-                if hasattr(possible_self_or_cls, "__class__"):
-                    if isinstance(possible_self_or_cls, type):
-                        # classmethod
-                        class_name = possible_self_or_cls.__name__
-                    else:
-                        # instance method
-                        class_name = possible_self_or_cls.__class__.__name__
+            class_name = self._extract_class_name(args)
 
             try:
                 result = func(*args, **kwargs)
                 duration = (time.perf_counter() - start) * 1000
+                log_data = self._build_log_data(
+                    func, class_name, duration, args, kwargs
+                )
+
                 if duration > self.config.slow_log_threshold_ms:
                     logger.log(
                         self.config.slow_log_level,
-                        f"{class_name}.{func.__name__} -> Slow execution: took %.2fms",
+                        "%s%s -> Slow execution: took %.2fms",
+                        f"{class_name}." if class_name else "",
+                        func.__name__,
                         duration,
+                        extra=log_data,
                     )
-                    # logger.warning("Slow execution: %s took %.2fms", func.__name__, duration)
                 else:
                     logger.info(
                         "Function %s%s executed successfully in %.2fms",
                         f"{class_name}." if class_name else "",
                         func.__name__,
                         duration,
-                        extra={
-                            "function_name": func.__name__,
-                            "class_name": class_name,
-                            "duration_ms": duration,
-                            "func_args": args,
-                            "func_kwargs": kwargs,
-                        },
+                        extra=log_data,
                     )
+
+                self._handle_export(log_data)
                 return result
+
             except Exception as e:
                 duration = (time.perf_counter() - start) * 1000
+                log_data = self._build_log_data(
+                    func, class_name, duration, args, kwargs, error=e
+                )
                 logger.error(
                     "Function %s%s failed after %.2fms: %s",
                     f"{class_name}." if class_name else "",
@@ -88,15 +158,9 @@ class TrackQuery(Generic[T]):
                     duration,
                     str(e),
                     exc_info=True,
-                    extra={
-                        "function_name": func.__name__,
-                        "class_name": class_name,
-                        "duration_ms": duration,
-                        "func_args": args,
-                        "func_kwargs": kwargs,
-                        "error": str(e),
-                    },
+                    extra=log_data,
                 )
-                raise
+                self._handle_export(log_data)
+                return None
 
         return update_wrapper(wrapped, func)

pyquerytracker/db/models.py

@@ -0,0 +1,20 @@
+from datetime import datetime
+
+from sqlalchemy import Column, DateTime, Float, Integer, String
+from sqlalchemy.orm import declarative_base
+
+Base = declarative_base()
+
+
+class TrackedQuery(Base):
+    __tablename__ = "tracked_queries"
+
+    id = Column(Integer, primary_key=True, index=True)
+    function_name = Column(String)
+    class_name = Column(String, nullable=True)
+    duration_ms = Column(Float)
+    timestamp = Column(DateTime, default=datetime.utcnow())
+    event = Column(String)  # "slow_execution", "normal_execution", "error"
+    func_args = Column(String)
+    func_kwargs = Column(String)
+    error = Column(String, nullable=True)

pyquerytracker/db/session.py

@@ -0,0 +1,8 @@
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+engine = create_engine(
+    "sqlite:///pyquerytracker/db/querytracker.db",
+    connect_args={"check_same_thread": False},
+)
+SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

pyquerytracker/db/writer.py

@@ -0,0 +1,64 @@
+from datetime import datetime, timedelta, timezone
+
+from sqlalchemy import select
+from sqlalchemy.exc import SQLAlchemyError
+
+from pyquerytracker.db.models import TrackedQuery
+from pyquerytracker.db.session import SessionLocal
+
+
+class DBWriter:
+    @staticmethod
+    def save(log_data: dict):
+        session = SessionLocal()
+        try:
+            entry = TrackedQuery(
+                function_name=log_data.get("function_name"),
+                class_name=log_data.get("class_name"),
+                duration_ms=log_data.get("duration_ms"),
+                event=log_data.get("event"),
+                func_args=log_data.get("func_args"),
+                func_kwargs=log_data.get("func_kwargs"),
+                error=log_data.get("error"),
+                timestamp=log_data.get("timestamp")
+                or datetime.now(timezone.utc),  # Ensure timestamp is set
+            )
+            session.add(entry)
+            session.commit()
+        except SQLAlchemyError as e:
+            session.rollback()
+            print(f"DBWriter error: {e}")
+        finally:
+            session.close()
+
+    @staticmethod
+    def fetch_all(minutes: int = 5):
+        session = SessionLocal()
+        try:
+            cutoff = datetime.utcnow() - timedelta(minutes=minutes)
+            stmt = (
+                select(TrackedQuery)
+                .where(TrackedQuery.timestamp >= cutoff)
+                .order_by(TrackedQuery.timestamp.desc())
+            )
+            results = session.execute(stmt).scalars().all()
+
+            return [
+                {
+                    "function_name": row.function_name,
+                    "class_name": row.class_name,
+                    "duration_ms": row.duration_ms,
+                    "timestamp": row.timestamp,
+                    "event": row.event,
+                    "error": row.error,
+                    "func_args": row.func_args,
+                    "func_kwargs": row.func_kwargs,
+                }
+                for row in results
+            ]
+
+        except SQLAlchemyError as e:
+            print(f"DB fetch error: {e}")
+            return []
+        finally:
+            session.close()

pyquerytracker/exporter/base.py

@@ -0,0 +1,25 @@
+from abc import ABC, abstractmethod
+
+from pyquerytracker.config import Config
+
+
+class Exporter(ABC):
+    def __init__(self, config: Config):
+        self.config = config
+
+    @abstractmethod
+    def append(self, data: dict) -> None:
+        pass
+
+    @abstractmethod
+    def flush(self) -> None:
+        pass
+
+
+class NullExporter:
+
+    def append(self, log_data):
+        pass
+
+    def flush(self):
+        pass

pyquerytracker/exporter/csv_exporter.py

@@ -0,0 +1,52 @@
+import atexit
+import csv
+import os
+from threading import Lock
+
+from pyquerytracker.exporter.base import Exporter
+from pyquerytracker.utils.logger import QueryLogger
+
+logger = QueryLogger.get_logger()
+
+
+class CsvExporter(Exporter):
+    def __init__(self, config):
+        super().__init__(config)
+        self._lock = Lock()
+        self._buffer = []
+        self._header_written = os.path.exists(self.config.export_path)
+        # Ensure logs are flushed when program exits.
+        atexit.register(self.flush)
+
+    def append(self, data: dict):
+        with self._lock:
+            self._buffer.append(data)
+
+    def flush(self):
+
+        with self._lock:
+            if not self._buffer:
+                return
+
+            os.makedirs(os.path.dirname(self.config.export_path), exist_ok=True)
+
+            # Gather all possible fieldnames (union of keys)
+            all_keys = set()
+            for entry in self._buffer:
+                all_keys.update(entry.keys())
+            fieldnames = sorted(all_keys)  # consistent ordering
+
+            with open(self.config.export_path, "a", newline="", encoding="utf-8") as f:
+                writer = csv.DictWriter(f, fieldnames=fieldnames)
+
+                if not self._header_written:
+                    writer.writeheader()
+                    self._header_written = True
+
+                for row in self._buffer:
+                    # Fill missing keys with None
+                    full_row = {key: row.get(key) for key in fieldnames}
+                    writer.writerow(full_row)
+
+                logger.info("Flushed %d logs to CSV", len(self._buffer))
+                self._buffer.clear()

pyquerytracker/exporter/json_exporter.py

@@ -0,0 +1,47 @@
+import atexit
+import json
+import os
+from threading import Lock
+
+from pyquerytracker.exporter.base import Exporter
+from pyquerytracker.utils.logger import QueryLogger
+
+logger = QueryLogger.get_logger()
+
+
+class JsonExporter(Exporter):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self._lock = Lock()
+        self._buffer = []
+        atexit.register(self.flush)
+
+    def append(self, data: dict):
+        with self._lock:
+            self._buffer.append(data)
+
+    def flush(self):
+        with self._lock:
+            if not self._buffer:
+                return
+
+            os.makedirs(os.path.dirname(self.config.export_path), exist_ok=True)
+
+            existing_data = []
+            if os.path.exists(self.config.export_path):
+                try:
+                    with open(self.config.export_path, "r", encoding="utf-8") as f:
+                        existing_data = json.load(f)
+                        if not isinstance(existing_data, list):
+                            existing_data = []
+                except json.JSONDecodeError:
+                    pass
+
+            existing_data.extend(self._buffer)
+
+            with open(self.config.export_path, "w", encoding="utf-8") as f:
+                json.dump(existing_data, f, indent=2, default=str)
+
+            logger.info(f"Flushed {len(self._buffer)} logs to JSON")
+            self._buffer.clear()

pyquerytracker/exporter/manager.py

@@ -0,0 +1,32 @@
+from pyquerytracker.config import Config, ExportType
+from pyquerytracker.exporter.base import Exporter
+from pyquerytracker.exporter.csv_exporter import CsvExporter
+from pyquerytracker.exporter.json_exporter import JsonExporter
+
+
+class ExporterManager:
+    _exporter: Exporter = None
+
+    _exporter_classes = {
+        ExportType.CSV: CsvExporter,
+        ExportType.JSON: JsonExporter,
+    }
+
+    @classmethod
+    def create_exporter(cls, config: Config) -> Exporter:
+        exporter_cls = cls._exporter_classes.get(config.export_type)
+        if not exporter_cls:
+            raise ValueError(f"Unsupported export type: {config.export_type}")
+        exporter = exporter_cls(config)
+        cls.set(exporter)
+        return exporter
+
+    @staticmethod
+    def set(exporter: Exporter):
+        ExporterManager._exporter = exporter
+
+    @staticmethod
+    def get() -> Exporter:
+        if ExporterManager._exporter is None:
+            raise RuntimeError("Exporter not set")
+        return ExporterManager._exporter

pyquerytracker/main.py

@@ -0,0 +1,6 @@
+if __name__ == "__main__":
+    import uvicorn
+
+    from pyquerytracker.api import app
+
+    uvicorn.run(app, host="127.0.0.1", port=8000)

pyquerytracker/tracker.py

@@ -0,0 +1,17 @@
+from datetime import datetime, timedelta
+from typing import Any, Dict, List
+
+# In-memory store to collect tracked query data
+query_data_store: List[Dict[str, Any]] = []
+
+
+def store_tracked_query(log: Dict[str, Any]):
+    """Store a single tracked query log entry."""
+    log["timestamp"] = datetime.utcnow()
+    query_data_store.append(log)
+
+
+def get_tracked_queries(minutes: int) -> List[Dict[str, Any]]:
+    """Return all tracked queries within the last `minutes`."""
+    cutoff = datetime.utcnow() - timedelta(minutes=minutes)
+    return [log for log in query_data_store if log["timestamp"] >= cutoff]

pyquerytracker/utils/logger.py

@@ -0,0 +1,18 @@
+import logging
+
+
+class QueryLogger:
+    @staticmethod
+    def get_logger(name: str = "pyquerytracker") -> logging.Logger:
+        logger = logging.getLogger(name)
+
+        if not logger.handlers:
+            handler = logging.StreamHandler()
+            fmt = logging.Formatter(
+                "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+            )
+            handler.setFormatter(fmt)
+            logger.addHandler(handler)
+            logger.setLevel(logging.INFO)
+
+        return logger

pyquerytracker/websocket.py

@@ -0,0 +1,33 @@
+import asyncio
+from typing import List
+
+from fastapi import WebSocket, WebSocketDisconnect
+
+from pyquerytracker.db.writer import DBWriter
+
+connected_clients: List[WebSocket] = []
+
+
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    connected_clients.append(websocket)
+    try:
+        while True:
+            await asyncio.sleep(2)  # every 2 seconds
+            recent_logs = DBWriter.fetch_all(minutes=5)  # or a custom method
+            await websocket.send_json(recent_logs)
+    except WebSocketDisconnect:
+        pass
+    finally:
+        connected_clients.remove(websocket)
+
+
+async def broadcast(message: str):
+    disconnected = []
+    for client in connected_clients:
+        try:
+            await client.send_text(message)
+        except Exception:
+            disconnected.append(client)
+    for client in disconnected:
+        connected_clients.remove(client)