feat(tasks): auto-capture worker log output for all tasks via TaskLogCollector
Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
This commit is contained in:
+69
-1
@@ -1,19 +1,87 @@
|
|||||||
|
import logging
|
||||||
|
import threading
|
||||||
|
from collections import defaultdict
|
||||||
|
|
||||||
from app.database import SessionLocal
|
from app.database import SessionLocal
|
||||||
from app.models import ProcessingLog
|
from app.models import ProcessingLog
|
||||||
|
|
||||||
|
|
||||||
|
class TaskLogCollector(logging.Handler):
|
||||||
|
"""
|
||||||
|
A logging handler that buffers log messages per Celery task ID.
|
||||||
|
|
||||||
|
When log_task_progress() is called, it drains the buffered messages
|
||||||
|
for that task and stores them in the ProcessingLog.detail field.
|
||||||
|
This captures all logger.info/error/warning output automatically.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super().__init__()
|
||||||
|
self._buffers = defaultdict(list)
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
|
||||||
|
def emit(self, record: logging.LogRecord) -> None:
|
||||||
|
"""Buffer a log record if it contains a task ID marker like [task-id]."""
|
||||||
|
try:
|
||||||
|
msg = self.format(record)
|
||||||
|
# Extract task_id from messages formatted as "[task_id] ..."
|
||||||
|
if msg and "[" in msg:
|
||||||
|
start = msg.index("[")
|
||||||
|
end = msg.index("]", start)
|
||||||
|
task_id = msg[start + 1 : end].strip()
|
||||||
|
if task_id and len(task_id) >= 8:
|
||||||
|
with self._lock:
|
||||||
|
self._buffers[task_id].append(msg)
|
||||||
|
except (ValueError, IndexError):
|
||||||
|
pass
|
||||||
|
|
||||||
|
def drain(self, task_id: str) -> str:
|
||||||
|
"""Return and clear all buffered messages for a task ID."""
|
||||||
|
with self._lock:
|
||||||
|
messages = self._buffers.pop(task_id, [])
|
||||||
|
return "\n".join(messages) if messages else ""
|
||||||
|
|
||||||
|
|
||||||
|
# Singleton collector instance
|
||||||
|
_collector = TaskLogCollector()
|
||||||
|
_collector.setLevel(logging.DEBUG)
|
||||||
|
_collector_installed = False
|
||||||
|
|
||||||
|
|
||||||
|
def _ensure_collector_installed() -> None:
|
||||||
|
"""Install the TaskLogCollector on the root logger (once)."""
|
||||||
|
global _collector_installed
|
||||||
|
if not _collector_installed:
|
||||||
|
root = logging.getLogger()
|
||||||
|
# Avoid duplicate handlers
|
||||||
|
if _collector not in root.handlers:
|
||||||
|
root.addHandler(_collector)
|
||||||
|
_collector_installed = True
|
||||||
|
|
||||||
|
|
||||||
def log_task_progress(task_id, step_name, status, message=None, file_id=None, detail=None):
|
def log_task_progress(task_id, step_name, status, message=None, file_id=None, detail=None):
|
||||||
"""
|
"""
|
||||||
Logs the progress of a Celery task to the database.
|
Logs the progress of a Celery task to the database.
|
||||||
|
|
||||||
|
If no explicit detail is provided, automatically drains any buffered
|
||||||
|
worker log output for this task ID and stores it as the detail.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
task_id: The Celery task ID
|
task_id: The Celery task ID
|
||||||
step_name: Name of the processing step
|
step_name: Name of the processing step
|
||||||
status: Current status (pending, in_progress, success, failure)
|
status: Current status (pending, in_progress, success, failure)
|
||||||
message: Short summary message
|
message: Short summary message
|
||||||
file_id: Optional associated file record ID
|
file_id: Optional associated file record ID
|
||||||
detail: Optional verbose log output for diagnostics
|
detail: Optional verbose log output for diagnostics.
|
||||||
|
If not provided, buffered logger output is used automatically.
|
||||||
"""
|
"""
|
||||||
|
# Auto-capture buffered log output when no explicit detail is given
|
||||||
|
if detail is None and task_id:
|
||||||
|
_ensure_collector_installed()
|
||||||
|
collected = _collector.drain(task_id)
|
||||||
|
if collected:
|
||||||
|
detail = collected
|
||||||
|
|
||||||
with SessionLocal() as db:
|
with SessionLocal() as db:
|
||||||
log_entry = ProcessingLog(
|
log_entry = ProcessingLog(
|
||||||
task_id=task_id,
|
task_id=task_id,
|
||||||
|
|||||||
@@ -4,6 +4,8 @@ Tests for app/utils/logging.py
|
|||||||
Tests task progress logging functionality.
|
Tests task progress logging functionality.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
from unittest.mock import Mock, patch, MagicMock
|
from unittest.mock import Mock, patch, MagicMock
|
||||||
|
|
||||||
@@ -176,3 +178,82 @@ class TestTaskLogging:
|
|||||||
assert mock_processing_log.called
|
assert mock_processing_log.called
|
||||||
mock_db.add.assert_called()
|
mock_db.add.assert_called()
|
||||||
mock_db.commit.assert_called()
|
mock_db.commit.assert_called()
|
||||||
|
|
||||||
|
@patch("app.utils.logging.SessionLocal")
|
||||||
|
@patch("app.utils.logging.ProcessingLog")
|
||||||
|
def test_log_task_progress_with_explicit_detail(self, mock_processing_log, mock_session_local):
|
||||||
|
"""Test logging with explicit detail preserves it."""
|
||||||
|
from app.utils.logging import log_task_progress
|
||||||
|
|
||||||
|
mock_db = MagicMock()
|
||||||
|
mock_session_local.return_value.__enter__.return_value = mock_db
|
||||||
|
|
||||||
|
mock_log_entry = Mock()
|
||||||
|
mock_processing_log.return_value = mock_log_entry
|
||||||
|
|
||||||
|
log_task_progress(
|
||||||
|
task_id="task-explicit",
|
||||||
|
step_name="test_step",
|
||||||
|
status="success",
|
||||||
|
message="Short message",
|
||||||
|
detail="Verbose detail output",
|
||||||
|
)
|
||||||
|
|
||||||
|
mock_processing_log.assert_called_once_with(
|
||||||
|
task_id="task-explicit",
|
||||||
|
step_name="test_step",
|
||||||
|
status="success",
|
||||||
|
message="Short message",
|
||||||
|
file_id=None,
|
||||||
|
detail="Verbose detail output",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.unit
|
||||||
|
class TestTaskLogCollector:
|
||||||
|
"""Test the TaskLogCollector handler."""
|
||||||
|
|
||||||
|
def test_collector_buffers_log_messages(self):
|
||||||
|
"""Test that the collector buffers messages by task ID."""
|
||||||
|
from app.utils.logging import TaskLogCollector
|
||||||
|
|
||||||
|
collector = TaskLogCollector()
|
||||||
|
collector.setFormatter(logging.Formatter("%(message)s"))
|
||||||
|
|
||||||
|
logger = logging.getLogger("test_collector")
|
||||||
|
logger.addHandler(collector)
|
||||||
|
logger.setLevel(logging.DEBUG)
|
||||||
|
|
||||||
|
logger.info("[abc12345-task] Step 1 starting")
|
||||||
|
logger.info("[abc12345-task] Step 1 complete")
|
||||||
|
logger.info("[other-task-id] Different task")
|
||||||
|
|
||||||
|
result = collector.drain("abc12345-task")
|
||||||
|
assert "Step 1 starting" in result
|
||||||
|
assert "Step 1 complete" in result
|
||||||
|
assert "Different task" not in result
|
||||||
|
|
||||||
|
# After drain, buffer should be empty
|
||||||
|
assert collector.drain("abc12345-task") == ""
|
||||||
|
|
||||||
|
# Other task still has its messages
|
||||||
|
result2 = collector.drain("other-task-id")
|
||||||
|
assert "Different task" in result2
|
||||||
|
|
||||||
|
logger.removeHandler(collector)
|
||||||
|
|
||||||
|
def test_collector_ignores_short_ids(self):
|
||||||
|
"""Test that the collector ignores short bracketed strings."""
|
||||||
|
from app.utils.logging import TaskLogCollector
|
||||||
|
|
||||||
|
collector = TaskLogCollector()
|
||||||
|
collector.setFormatter(logging.Formatter("%(message)s"))
|
||||||
|
|
||||||
|
logger = logging.getLogger("test_short_ids")
|
||||||
|
logger.addHandler(collector)
|
||||||
|
logger.setLevel(logging.DEBUG)
|
||||||
|
|
||||||
|
logger.info("[OK] short id")
|
||||||
|
assert collector.drain("OK") == ""
|
||||||
|
|
||||||
|
logger.removeHandler(collector)
|
||||||
|
|||||||
Reference in New Issue
Block a user