Feat: log handler (#2)

* add log handler that writes to questdb

* remove accidental pycache commits

Author: idatsy

.gitignore

@@ -2,3 +2,4 @@
 poetry.lock
 dist/
 **/__pycache__/
+*.pyc

py_questdb/__init__.py

@@ -1,3 +1,3 @@
 from py_questdb.db import QuestDB
 
-__all__ = ['QuestDB']
+__all__ = ["QuestDB"]

py_questdb/__pycache__/__init__.cpython-312.pyc

@@ -37,6 +37,10 @@ def __init__(self, host: str = "localhost", port: int = 9000, username: str | No
         self.client.establish()  # unclear if this blocks, API keeps changing!
         self.session = None
 
+    def flush(self) -> None:
+        """Flush the QuestDB client buffer."""
+        self.client.flush()
+
     def write(self, *messages: QuestMessage) -> None:
         """
         Write multiple messages to QuestDB.
@@ -125,9 +129,7 @@ def parse_query_response(response: bytes) -> QuestDBResponse:
         return msgspec.json.decode(response, type=QuestDBResponse)
 
     @staticmethod
-    def parse_and_yield_query_response(
-        response: QuestDBResponse, into_type: Type[T] | None
-    ) -> Iterable[T] | Iterable[dict]:
+    def parse_and_yield_query_response(response: QuestDBResponse, into_type: Type[T] | None) -> Iterable[T] | Iterable[dict]:
         """
         Parse and yield rows from a QuestDBResponse.
 
@@ -143,8 +145,7 @@ def parse_and_yield_query_response(
         for row in response.dataset:
             converted_row = {
                 field.name: converter(value) if value is not None else None
-                for field, converter, value
-                in zip(response.columns, type_converter, row)
+                for field, converter, value in zip(response.columns, type_converter, row)
             }
 
             yield into_type(**converted_row) if into_type else converted_row
@@ -238,6 +239,8 @@ def query_df_sync(self, query_string: str) -> pd.DataFrame:
 
     async def close(self):
         """Close the QuestDB connection and any open sessions."""
+        self.flush()
+
         if self.session:
             await self.session.close()
         self.client.close()
@@ -251,5 +254,7 @@ async def __aexit__(
         self, exc_type: Type[BaseException] | None, exc_val: BaseException | None, exc_tb: TracebackType | None
     ) -> None:
         """Async context manager exit."""
+        self.flush()
+
         if self.session:
             await self.close()

py_questdb/__pycache__/db.cpython-312.pyc

@@ -0,0 +1,87 @@
+"""
+USAGE:
+    questdb_client = QuestDB()
+    handler = QuestDBLogHandler(questdb_client=questdb_client)
+    logger.addHandler(handler)
+"""
+
+import logging
+from logging.handlers import QueueHandler, QueueListener
+from multiprocessing import Queue
+from datetime import datetime
+from typing import Optional
+from py_questdb import QuestDB
+
+
+class QuestDBLogHandler(QueueHandler):
+    def __init__(self, questdb_client: QuestDB, table_name: str = "_logs", level: Optional[str] = None):
+        self.queue = Queue()
+        super().__init__(self.queue)
+
+        if level:
+            self.setLevel(level)
+
+        self.inner_handler = InnerQuestDBLogHandler(questdb_client=questdb_client, table_name=table_name)
+
+        self.listener = QueueListener(self.queue, self.inner_handler, respect_handler_level=True)
+        self.listener.start()
+
+    def prepare(self, record: logging.LogRecord) -> logging.LogRecord:
+        if record.exc_info:
+            record.exc_info = logging._defaultFormatter.formatException(record.exc_info)  # noqa
+        return record
+
+    def close(self) -> None:
+        try:
+            if self.listener:
+                self.listener.stop()
+                self.listener = None
+        except Exception as e:
+            print(f"Error stopping QueueListener: {e}")
+        finally:
+            self.inner_handler.questdb_client.flush()
+            super().close()
+
+    def __del__(self):
+        self.close()
+
+
+class InnerQuestDBLogHandler(logging.Handler):
+    def __init__(self, questdb_client: QuestDB, table_name: str = "_logs"):
+        super().__init__()
+        self.questdb_client = questdb_client
+        self.table_name = table_name
+
+    def emit(self, record: logging.LogRecord):
+        try:
+            print(f"Received {record=}")
+            log_message = self.format_record(record)
+            print(f"Writing {log_message=}")
+            self.questdb_client.client.row(**log_message)
+        except Exception as e:
+            print(f"Failed to write log to QuestDB: {e}")
+
+    def format_record(self, record: logging.LogRecord) -> dict:
+        print(record.stack_info)
+        return {
+            "table_name": self.table_name,
+            "symbols": {
+                "level": record.levelname,
+                "logger": record.name,
+                "filename": record.filename,
+                "funcName": record.funcName,
+                "module": record.module,
+                "processName": record.processName,
+                "threadName": record.threadName,
+                "pathname": record.pathname,
+            },
+            "columns": {
+                "message": record.getMessage(),
+                "lineno": record.lineno,
+                "process": record.process,
+                "thread": record.thread,
+                "exc_info": record.exc_info,
+                "stack_info": record.stack_info,
+            },
+            "at": datetime.fromtimestamp(record.created),
+        }

py_questdb/__pycache__/db_types.cpython-312.pyc

@@ -0,0 +1,93 @@
+import pytest
+import logging
+from unittest.mock import MagicMock, patch
+from py_questdb import QuestDB
+from py_questdb.log_handler import QuestDBLogHandler, InnerQuestDBLogHandler
+
+
+def setup_logging(questdb_client: QuestDB):
+    handler = QuestDBLogHandler(questdb_client)
+    root_logger = logging.getLogger()
+    root_logger.setLevel(logging.DEBUG)
+    root_logger.addHandler(handler)
+    return handler
+
+
+def worker_process(worker_id):
+    logger = logging.getLogger(f"worker_{worker_id}")
+    if worker_id == 2:
+        try:
+            raise ValueError(f"This is an exception from worker {worker_id}")
+        except ValueError as e:
+            logger.error(f"This is an error message from worker {worker_id} {e}", exc_info=True, stack_info=True)
+    logger.info(f"This is an info message from worker {worker_id}")
+
+
+@pytest.fixture
+def mock_questdb_client():
+    mock_client = MagicMock()
+    mock_questdb = MagicMock(spec=QuestDB)
+    mock_questdb.client = mock_client
+    return mock_questdb
+
+
+@pytest.fixture
+def questdb_log_handler(mock_questdb_client):
+    return QuestDBLogHandler(mock_questdb_client)
+
+
+def test_questdb_log_handler_init(mock_questdb_client):
+    handler = QuestDBLogHandler(mock_questdb_client)
+    assert isinstance(handler.inner_handler, InnerQuestDBLogHandler)
+    assert handler.listener is not None
+
+
+def test_questdb_log_handler_prepare(questdb_log_handler):
+    record = logging.LogRecord(
+        name="test_logger", level=logging.INFO, pathname="test_path", lineno=42, msg="Test message", args=(), exc_info=None
+    )
+    prepared_record = questdb_log_handler.prepare(record)
+    assert prepared_record == record
+
+
+def test_questdb_log_handler_close(questdb_log_handler):
+    with patch.object(questdb_log_handler.listener, "stop") as mock_stop:
+        questdb_log_handler.close()
+        mock_stop.assert_called_once()
+    assert questdb_log_handler.listener is None
+
+
+def test_inner_questdb_log_handler_emit(mock_questdb_client):
+    handler = InnerQuestDBLogHandler(mock_questdb_client)
+    record = logging.LogRecord(
+        name="test_logger", level=logging.INFO, pathname="test_path", lineno=42, msg="Test message", args=(), exc_info=None
+    )
+    handler.emit(record)
+    mock_questdb_client.client.row.assert_called_once()
+
+
+def test_setup_logging(mock_questdb_client):
+    with patch("logging.getLogger") as mock_get_logger:
+        mock_root_logger = MagicMock()
+        mock_get_logger.return_value = mock_root_logger
+        handler = setup_logging(mock_questdb_client)
+        assert isinstance(handler, QuestDBLogHandler)
+        mock_root_logger.setLevel.assert_called_with(logging.DEBUG)
+        mock_root_logger.addHandler.assert_called_with(handler)
+
+
+@patch("logging.getLogger")
+def test_worker_process(mock_get_logger):
+    mock_logger = MagicMock()
+    mock_get_logger.return_value = mock_logger
+
+    # Test normal case
+    worker_process(1)
+    mock_logger.info.assert_called_with("This is an info message from worker 1")
+
+    # Test exception case
+    worker_process(2)
+    mock_logger.error.assert_called()
+    assert "This is an error message from worker 2" in mock_logger.error.call_args[0][0]
+    assert mock_logger.error.call_args[1]["exc_info"] is True
+    assert mock_logger.error.call_args[1]["stack_info"] is True