diff --git a/.pylintrc b/.pylintrc index 93378e7..8a2efa4 100644 --- a/.pylintrc +++ b/.pylintrc @@ -9,7 +9,8 @@ disable= too-few-public-methods, fixme, duplicate-code, - logging-fstring-interpolation + logging-fstring-interpolation, + broad-exception-caught [REPORTS] reports=no diff --git a/examples/exporter/csv_exporter_1.py b/examples/exporter/csv_exporter_1.py new file mode 100644 index 0000000..b9eddb3 --- /dev/null +++ b/examples/exporter/csv_exporter_1.py @@ -0,0 +1,40 @@ +import os +import time + +from pyquerytracker import TrackQuery +from pyquerytracker.config import configure, ExportType + +os.makedirs("logs-csv", exist_ok=True) + +configure( + slow_log_threshold_ms=50.0, + slow_log_level=20, # INFO + export_type=ExportType.CSV, + export_path="logs-csv/query_logs_2.csv", +) + + +@TrackQuery() +def process_data(x, y): + time.sleep(8) + return x + y + + +@TrackQuery() +def failing_task(): + time.sleep(0.03) + raise RuntimeError("This failed intentionally.") + + +# Run +print("Result:", process_data(5, 7)) +print("Result:", process_data(5, 7)) +print("Result:", process_data(5, 7)) +print("Result:", process_data(5, 7)) +print("Result:", process_data(5, 7)) +print("Result:", process_data(5, 7)) + +try: + failing_task() +except Exception: + pass diff --git a/pyquerytracker/config.py b/pyquerytracker/config.py index 74d6b42..6f1f391 100644 --- a/pyquerytracker/config.py +++ b/pyquerytracker/config.py @@ -35,6 +35,8 @@ class Config: # TODO: Adding export functionality slow_log_threshold_ms: float = 100.0 slow_log_level: int = logging.WARNING + export_type: Optional[ExportType] = None + export_path: Optional[str] = None _config: Config = Config() @@ -43,6 +45,8 @@ class Config: def configure( slow_log_threshold_ms: Optional[float] = None, slow_log_level: Optional[int] = None, + export_type: Optional[ExportType] = None, + export_path: Optional[str] = None, ): """ Configure global settings for query tracking. @@ -60,6 +64,10 @@ def configure( _config.slow_log_threshold_ms = slow_log_threshold_ms if slow_log_level is not None: _config.slow_log_level = slow_log_level + if export_type is not None: + _config.export_type = export_type + if export_path is not None: + _config.export_path = export_path def get_config() -> Config: diff --git a/pyquerytracker/core.py b/pyquerytracker/core.py index baf936d..9569cb7 100644 --- a/pyquerytracker/core.py +++ b/pyquerytracker/core.py @@ -1,19 +1,12 @@ import time -import logging from functools import update_wrapper from typing import Any, Callable, TypeVar, Generic from pyquerytracker.config import get_config +from pyquerytracker.utils.logger import QueryLogger +from pyquerytracker.exporter.manager import ExporterManager +from pyquerytracker.exporter.base import NullExporter -# 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") @@ -37,6 +30,12 @@ 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 __call__(self, func: Callable[..., T]) -> Callable[..., T]: def wrapped(*args: Any, **kwargs: Any) -> T: @@ -48,20 +47,31 @@ def wrapped(*args: Any, **kwargs: Any) -> T: 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__ try: result = func(*args, **kwargs) duration = (time.perf_counter() - start) * 1000 + log_data = { + "event": ( + "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 duration > self.config.slow_log_threshold_ms: logger.log( self.config.slow_log_level, - f"{class_name}.{func.__name__} -> Slow execution: took %.2fms", - duration, + f"{class_name}.{func.__name__} -> " + f"Slow execution: took {duration:.2f}ms", ) else: @@ -70,17 +80,23 @@ def wrapped(*args: Any, **kwargs: Any) -> T: 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.exporter.append(log_data) return result + except Exception as e: duration = (time.perf_counter() - start) * 1000 + log_data = { + "event": "error", + "function_name": func.__name__, + "class_name": class_name, + "duration_ms": duration, + "func_args": repr(args), + "func_kwargs": repr(kwargs), + "error": str(e), + } + logger.error( "Function %s%s failed after %.2fms: %s", f"{class_name}." if class_name else "", @@ -88,15 +104,11 @@ def wrapped(*args: Any, **kwargs: Any) -> 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.exporter.append(log_data) + # Exceptions are handled internally by the decorator, + # allowing the program to proceed smoothly + return None return update_wrapper(wrapped, func) diff --git a/pyquerytracker/exporter/__init__.py b/pyquerytracker/exporter/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/pyquerytracker/exporter/base.py b/pyquerytracker/exporter/base.py new file mode 100644 index 0000000..fce8efd --- /dev/null +++ b/pyquerytracker/exporter/base.py @@ -0,0 +1,24 @@ +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 diff --git a/pyquerytracker/exporter/csv_exporter.py b/pyquerytracker/exporter/csv_exporter.py new file mode 100644 index 0000000..ad0f090 --- /dev/null +++ b/pyquerytracker/exporter/csv_exporter.py @@ -0,0 +1,52 @@ +import atexit +import csv +import os +from threading import Lock + +from pyquerytracker.utils.logger import QueryLogger +from pyquerytracker.exporter.base import Exporter + +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() diff --git a/pyquerytracker/exporter/manager.py b/pyquerytracker/exporter/manager.py new file mode 100644 index 0000000..226e05d --- /dev/null +++ b/pyquerytracker/exporter/manager.py @@ -0,0 +1,25 @@ +from pyquerytracker.config import Config, ExportType +from pyquerytracker.exporter.csv_exporter import CsvExporter +from pyquerytracker.exporter.base import Exporter + + +class ExporterManager: + _exporter: Exporter = None + + @staticmethod + def create_exporter(config: Config) -> Exporter: + if config.export_type == ExportType.CSV: + exporter = CsvExporter(config) + ExporterManager.set(exporter) + return exporter + raise ValueError("Unsupported export type") + + @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 diff --git a/pyquerytracker/utils/logger.py b/pyquerytracker/utils/logger.py new file mode 100644 index 0000000..59e8b90 --- /dev/null +++ b/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 diff --git a/tests/test_core.py b/tests/test_core.py index e6e086f..338cf1b 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -42,8 +42,8 @@ def test_tracking_output_with_error(caplog): def failing_query(): raise ValueError("Test error") - with pytest.raises(ValueError): - failing_query() + result = failing_query() + assert result is None # Check the log records assert len(caplog.records) == 1