avtomatika 1.0b7__py3-none-any.whl → 1.0b9__py3-none-any.whl

avtomatika/engine.py

@@ -19,10 +19,12 @@ from .app_keys import (
     HTTP_SESSION_KEY,
     REPUTATION_CALCULATOR_KEY,
     REPUTATION_CALCULATOR_TASK_KEY,
+    S3_SERVICE_KEY,
     SCHEDULER_KEY,
     SCHEDULER_TASK_KEY,
     WATCHER_KEY,
     WATCHER_TASK_KEY,
+    WORKER_SERVICE_KEY,
     WS_MANAGER_KEY,
 )
 from .blueprint import StateMachineBlueprint
@@ -37,7 +39,9 @@ from .history.base import HistoryStorageBase
 from .history.noop import NoOpHistoryStorage
 from .logging_config import setup_logging
 from .reputation import ReputationCalculator
+from .s3 import S3Service
 from .scheduler import Scheduler
+from .services.worker_service import WorkerService
 from .storage.base import StorageBackend
 from .telemetry import setup_telemetry
 from .utils.webhook_sender import WebhookPayload, WebhookSender
@@ -54,7 +58,7 @@ def json_dumps(obj: Any) -> str:
     return dumps(obj).decode("utf-8")
 
 
-def json_response(data: Any, **kwargs: Any) -> web.Response:
+def json_response(data, **kwargs: Any) -> web.Response:
     return web.json_response(data, dumps=json_dumps, **kwargs)
 
 
@@ -69,8 +73,13 @@ class OrchestratorEngine:
         self.ws_manager = WebSocketManager()
         self.app = web.Application(middlewares=[compression_middleware])
         self.app[ENGINE_KEY] = self
+        self.worker_service = None
         self._setup_done = False
 
+        from rxon import HttpListener
+
+        self.rxon_listener = HttpListener(self.app)
+
     def register_blueprint(self, blueprint: StateMachineBlueprint) -> None:
         if self._setup_done:
             raise RuntimeError("Cannot register blueprints after engine setup.")
@@ -140,7 +149,75 @@ class OrchestratorEngine:
                 )
                 self.history_storage = NoOpHistoryStorage()
 
+    async def handle_rxon_message(self, message_type: str, payload: Any, context: dict) -> Any:
+        """Core handler for RXON protocol messages via any listener."""
+        from rxon.security import extract_cert_identity
+
+        from .security import verify_worker_auth
+
+        request = context.get("raw_request")
+        token = context.get("token")
+        cert_identity = extract_cert_identity(request) if request else None
+
+        worker_id_hint = context.get("worker_id_hint")
+
+        if not worker_id_hint:
+            if message_type == "poll" and isinstance(payload, str):
+                worker_id_hint = payload
+            elif isinstance(payload, dict) and "worker_id" in payload:
+                worker_id_hint = payload["worker_id"]
+            elif hasattr(payload, "worker_id"):
+                worker_id_hint = payload.worker_id
+
+        try:
+            auth_worker_id = await verify_worker_auth(self.storage, self.config, token, cert_identity, worker_id_hint)
+        except PermissionError as e:
+            raise web.HTTPUnauthorized(text=str(e)) from e
+        except ValueError as e:
+            raise web.HTTPBadRequest(text=str(e)) from e
+
+        if message_type == "register":
+            return await self.worker_service.register_worker(payload)
+
+        elif message_type == "poll":
+            return await self.worker_service.get_next_task(auth_worker_id)
+
+        elif message_type == "result":
+            return await self.worker_service.process_task_result(payload, auth_worker_id)
+
+        elif message_type == "heartbeat":
+            return await self.worker_service.update_worker_heartbeat(auth_worker_id, payload)
+
+        elif message_type == "sts_token":
+            if cert_identity is None:
+                raise web.HTTPForbidden(text="Unauthorized: mTLS certificate required to issue access token.")
+            return await self.worker_service.issue_access_token(auth_worker_id)
+
+        elif message_type == "websocket":
+            ws = payload
+            await self.ws_manager.register(auth_worker_id, ws)
+            try:
+                from aiohttp import WSMsgType
+
+                async for msg in ws:
+                    if msg.type == WSMsgType.TEXT:
+                        try:
+                            data = msg.json()
+                            await self.ws_manager.handle_message(auth_worker_id, data)
+                        except Exception as e:
+                            logger.error(f"Error processing WebSocket message from {auth_worker_id}: {e}")
+                    elif msg.type == WSMsgType.ERROR:
+                        break
+            finally:
+                await self.ws_manager.unregister(auth_worker_id)
+            return None
+
     async def on_startup(self, app: web.Application) -> None:
+        # Fail Fast: Check Storage Connection
+        if not await self.storage.ping():
+            logger.critical("Failed to connect to Storage Backend (Redis). Exiting.")
+            raise RuntimeError("Storage Backend is unavailable.")
+
         try:
             from opentelemetry.instrumentation.aiohttp_client import (
                 AioHttpClientInstrumentor,
@@ -152,6 +229,8 @@ class OrchestratorEngine:
                 "opentelemetry-instrumentation-aiohttp-client not found. AIOHTTP client instrumentation is disabled."
             )
         await self._setup_history_storage()
+        # Start history background worker
+        await self.history_storage.start()
 
         # Load client configs if the path is provided
         if self.config.CLIENTS_CONFIG_PATH:
@@ -188,6 +267,7 @@ class OrchestratorEngine:
 
         app[HTTP_SESSION_KEY] = ClientSession()
         self.webhook_sender = WebhookSender(app[HTTP_SESSION_KEY])
+        self.webhook_sender.start()
         self.dispatcher = Dispatcher(self.storage, self.config)
         app[DISPATCHER_KEY] = self.dispatcher
         app[EXECUTOR_KEY] = JobExecutor(self, self.history_storage)
@@ -196,6 +276,10 @@ class OrchestratorEngine:
         app[HEALTH_CHECKER_KEY] = HealthChecker(self)
         app[SCHEDULER_KEY] = Scheduler(self)
         app[WS_MANAGER_KEY] = self.ws_manager
+        app[S3_SERVICE_KEY] = S3Service(self.config, self.history_storage)
+
+        self.worker_service = WorkerService(self.storage, self.history_storage, self.config, self)
+        app[WORKER_SERVICE_KEY] = self.worker_service
 
         app[EXECUTOR_TASK_KEY] = create_task(app[EXECUTOR_KEY].run())
         app[WATCHER_TASK_KEY] = create_task(app[WATCHER_KEY].run())
@@ -203,8 +287,12 @@ class OrchestratorEngine:
         app[HEALTH_CHECKER_TASK_KEY] = create_task(app[HEALTH_CHECKER_KEY].run())
         app[SCHEDULER_TASK_KEY] = create_task(app[SCHEDULER_KEY].run())
 
+        await self.rxon_listener.start(self.handle_rxon_message)
+
     async def on_shutdown(self, app: web.Application) -> None:
         logger.info("Shutdown sequence started.")
+        await self.rxon_listener.stop()
+
         app[EXECUTOR_KEY].stop()
         app[WATCHER_KEY].stop()
         app[REPUTATION_CALCULATOR_KEY].stop()
@@ -220,6 +308,13 @@ class OrchestratorEngine:
         logger.info("Closing WebSocket connections...")
         await self.ws_manager.close_all()
 
+        logger.info("Stopping WebhookSender...")
+        await self.webhook_sender.stop()
+
+        if S3_SERVICE_KEY in app:
+            logger.info("Closing S3 Service...")
+            await app[S3_SERVICE_KEY].close()
+
         logger.info("Cancelling background tasks...")
         app[HEALTH_CHECKER_TASK_KEY].cancel()
         app[WATCHER_TASK_KEY].cancel()
@@ -256,6 +351,7 @@ class OrchestratorEngine:
         blueprint_name: str,
         initial_data: dict[str, Any],
         source: str = "internal",
+        tracing_context: dict[str, str] | None = None,
     ) -> str:
         """Creates a job directly, bypassing the HTTP API layer.
         Useful for internal schedulers and triggers.
@@ -279,7 +375,7 @@ class OrchestratorEngine:
             "initial_data": initial_data,
             "state_history": {},
             "status": JOB_STATUS_PENDING,
-            "tracing_context": {},
+            "tracing_context": tracing_context or {},
             "client_config": client_config,
         }
         await self.storage.save_job_state(job_id, job_state)
@@ -352,23 +448,48 @@ class OrchestratorEngine:
         )
 
         # Run in background to not block the main flow
-        create_task(self.webhook_sender.send(webhook_url, payload))
+        await self.webhook_sender.send(webhook_url, payload)
 
     def run(self) -> None:
         self.setup()
+        ssl_context = None
+        if self.config.TLS_ENABLED:
+            from rxon.security import create_server_ssl_context
+
+            ssl_context = create_server_ssl_context(
+                cert_path=self.config.TLS_CERT_PATH,
+                key_path=self.config.TLS_KEY_PATH,
+                ca_path=self.config.TLS_CA_PATH,
+                require_client_cert=self.config.TLS_REQUIRE_CLIENT_CERT,
+            )
+            print(f"TLS enabled. mTLS required: {self.config.TLS_REQUIRE_CLIENT_CERT}")
+
         print(
             f"Starting OrchestratorEngine API server on {self.config.API_HOST}:{self.config.API_PORT} in blocking mode."
         )
-        web.run_app(self.app, host=self.config.API_HOST, port=self.config.API_PORT)
+        web.run_app(self.app, host=self.config.API_HOST, port=self.config.API_PORT, ssl_context=ssl_context)
 
     async def start(self):
         """Starts the orchestrator engine non-blockingly."""
         self.setup()
         self.runner = web.AppRunner(self.app)
         await self.runner.setup()
-        self.site = web.TCPSite(self.runner, self.config.API_HOST, self.config.API_PORT)
+
+        ssl_context = None
+        if self.config.TLS_ENABLED:
+            from rxon.security import create_server_ssl_context
+
+            ssl_context = create_server_ssl_context(
+                cert_path=self.config.TLS_CERT_PATH,
+                key_path=self.config.TLS_KEY_PATH,
+                ca_path=self.config.TLS_CA_PATH,
+                require_client_cert=self.config.TLS_REQUIRE_CLIENT_CERT,
+            )
+
+        self.site = web.TCPSite(self.runner, self.config.API_HOST, self.config.API_PORT, ssl_context=ssl_context)
         await self.site.start()
-        print(f"OrchestratorEngine API server running on http://{self.config.API_HOST}:{self.config.API_PORT}")
+        protocol = "https" if self.config.TLS_ENABLED else "http"
+        print(f"OrchestratorEngine API server running on {protocol}://{self.config.API_HOST}:{self.config.API_PORT}")
 
     async def stop(self):
         """Stops the orchestrator engine."""

avtomatika/executor.py

@@ -47,6 +47,17 @@ except ImportError:
     inject = NoOpPropagate().inject
     TraceContextTextMapPropagator = NoOpTraceContextTextMapPropagator()  # Instantiate the class
 
+from .app_keys import S3_SERVICE_KEY
+from .constants import (
+    JOB_STATUS_ERROR,
+    JOB_STATUS_FAILED,
+    JOB_STATUS_FINISHED,
+    JOB_STATUS_PENDING,
+    JOB_STATUS_QUARANTINED,
+    JOB_STATUS_RUNNING,
+    JOB_STATUS_WAITING_FOR_PARALLEL,
+    JOB_STATUS_WAITING_FOR_WORKER,
+)
 from .context import ActionFactory
 from .data_types import ClientConfig, JobContext
 from .history.base import HistoryStorageBase
@@ -58,7 +69,7 @@ logger = getLogger(
     __name__
 )  # Re-declare logger after potential redefinition in except block if opentelemetry was missing
 
-TERMINAL_STATES = {"finished", "failed", "error", "quarantined"}
+TERMINAL_STATES = {JOB_STATUS_FINISHED, JOB_STATUS_FAILED, JOB_STATUS_ERROR, JOB_STATUS_QUARANTINED}
 
 
 class JobExecutor:
@@ -74,7 +85,7 @@ class JobExecutor:
         self._running = False
         self._processing_messages: set[str] = set()
 
-    async def _process_job(self, job_id: str, message_id: str):
+    async def _process_job(self, job_id: str, message_id: str) -> None:
         """The core logic for processing a single job dequeued from storage."""
         if message_id in self._processing_messages:
             return
@@ -143,6 +154,11 @@ class JobExecutor:
                     plan=client_config_dict.get("plan", "unknown"),
                     params=client_config_dict.get("params", {}),
                 )
+
+                # Get TaskFiles if S3 service is available
+                s3_service = self.engine.app.get(S3_SERVICE_KEY)
+                task_files = s3_service.get_task_files(job_id) if s3_service else None
+
                 context = JobContext(
                     job_id=job_id,
                     current_state=job_state["current_state"],
@@ -153,6 +169,7 @@ class JobExecutor:
                     data_stores=SimpleNamespace(**blueprint.data_stores),
                     tracing_context=tracing_context,
                     aggregation_results=job_state.get("aggregation_results"),
+                    task_files=task_files,
                 )
 
                 try:
@@ -173,12 +190,17 @@ class JobExecutor:
                         params_to_inject["context"] = context
                         if "actions" in param_names:
                             params_to_inject["actions"] = action_factory
+                        if "task_files" in param_names:
+                            params_to_inject["task_files"] = task_files
                     else:
                         # New injection logic with prioritized lookup.
                         context_as_dict = context._asdict()
                         for param_name in param_names:
+                            # Direct injection of task_files
+                            if param_name == "task_files":
+                                params_to_inject[param_name] = task_files
                             # Look in JobContext fields first.
-                            if param_name in context_as_dict:
+                            elif param_name in context_as_dict:
                                 params_to_inject[param_name] = context_as_dict[param_name]
                             # Then look in state_history (data from previous steps/workers).
                             elif param_name in context.state_history:
@@ -251,19 +273,24 @@ class JobExecutor:
         # When transitioning to a new state, reset the retry counter.
         job_state["retry_count"] = 0
         job_state["current_state"] = next_state
-        job_state["status"] = "running"
+        job_state["status"] = JOB_STATUS_RUNNING
         await self.storage.save_job_state(job_id, job_state)
 
         if next_state not in TERMINAL_STATES:
             await self.storage.enqueue_job(job_id)
         else:
             logger.info(f"Job {job_id} reached terminal state {next_state}")
+
+            # Clean up S3 files if service is available
+            s3_service = self.engine.app.get(S3_SERVICE_KEY)
+            if s3_service:
+                task_files = s3_service.get_task_files(job_id)
+                if task_files:
+                    # Run cleanup in background to not block response
+                    create_task(task_files.cleanup())
+
             await self._check_and_resume_parent(job_state)
-            # Send webhook for finished/failed jobs
-            event_type = "job_finished" if next_state == "finished" else "job_failed"
-            # Since _check_and_resume_parent is for sub-jobs, we only send webhook if it's a top-level job
-            # or if the user explicitly requested it for sub-jobs (by providing webhook_url).
-            # The current logic stores webhook_url in job_state, so we just check it.
+            event_type = "job_finished" if next_state == JOB_STATUS_FINISHED else "job_failed"
             await self.engine.send_job_webhook(job_state, event_type)
 
     async def _handle_dispatch(
@@ -292,21 +319,15 @@ class JobExecutor:
             logger.info(f"Job {job_id} is now paused, awaiting human approval.")
         else:
             logger.info(f"Job {job_id} dispatching task: {task_info}")
-
             now = monotonic()
-            # Safely get timeout, falling back to the global config if not provided in the task.
-            # This prevents TypeErrors if 'timeout_seconds' is missing.
             timeout_seconds = task_info.get("timeout_seconds") or self.engine.config.WORKER_TIMEOUT_SECONDS
             timeout_at = now + timeout_seconds
-
-            # Set status to waiting and add to watch list *before* dispatching
-            job_state["status"] = "waiting_for_worker"
+            job_state["status"] = JOB_STATUS_WAITING_FOR_WORKER
             job_state["task_dispatched_at"] = now
             job_state["current_task_info"] = task_info  # Save for retries
             job_state["current_task_transitions"] = task_info.get("transitions", {})
             await self.storage.save_job_state(job_id, job_state)
             await self.storage.add_job_to_watch(job_id, timeout_at)
-
             await self.dispatcher.dispatch(job_state, task_info)
 
     async def _handle_run_blueprint(
@@ -334,7 +355,7 @@ class JobExecutor:
             "blueprint_name": sub_blueprint_info["blueprint_name"],
             "current_state": "start",
             "initial_data": sub_blueprint_info["initial_data"],
-            "status": "pending",
+            "status": JOB_STATUS_PENDING,
             "parent_job_id": parent_job_id,
         }
         await self.storage.save_job_state(child_job_id, child_job_state)
@@ -367,7 +388,7 @@ class JobExecutor:
         branch_task_ids = [str(uuid4()) for _ in tasks_to_dispatch]
 
         # Update job state for parallel execution
-        job_state["status"] = "waiting_for_parallel_tasks"
+        job_state["status"] = JOB_STATUS_WAITING_FOR_PARALLEL
         job_state["aggregation_target"] = aggregate_into
         job_state["active_branches"] = branch_task_ids
         job_state["aggregation_results"] = {}
@@ -445,7 +466,7 @@ class JobExecutor:
             logger.critical(
                 f"Job {job_id} has failed handler execution {max_retries + 1} times. Moving to quarantine.",
             )
-            job_state["status"] = "quarantined"
+            job_state["status"] = JOB_STATUS_QUARANTINED
             job_state["error_message"] = str(error)
             await self.storage.save_job_state(job_id, job_state)
             await self.storage.quarantine_job(job_id)
@@ -478,7 +499,7 @@ class JobExecutor:
             return
 
         # Determine the outcome of the child job to select the correct transition.
-        child_outcome = "success" if child_job_state["current_state"] == "finished" else "failure"
+        child_outcome = "success" if child_job_state["current_state"] == JOB_STATUS_FINISHED else "failure"
         transitions = parent_job_state.get("current_task_transitions", {})
         next_state = transitions.get(child_outcome, "failed")
 
@@ -493,7 +514,7 @@ class JobExecutor:
 
         # Update the parent job to its new state and re-enqueue it.
         parent_job_state["current_state"] = next_state
-        parent_job_state["status"] = "running"
+        parent_job_state["status"] = JOB_STATUS_RUNNING
         await self.storage.save_job_state(parent_job_id, parent_job_state)
         await self.storage.enqueue_job(parent_job_id)
 
@@ -522,7 +543,10 @@ class JobExecutor:
                 # Wait for an available slot before fetching a new job
                 await semaphore.acquire()
 
-                result = await self.storage.dequeue_job()
+                # Block for a configured time waiting for a job
+                block_time = self.engine.config.REDIS_STREAM_BLOCK_MS
+                result = await self.storage.dequeue_job(block=block_time if block_time > 0 else None)
+
                 if result:
                     job_id, message_id = result
                     task = create_task(self._process_job(job_id, message_id))
@@ -530,14 +554,18 @@ class JobExecutor:
                     # Release the semaphore slot when the task is done
                     task.add_done_callback(lambda _: semaphore.release())
                 else:
-                    # No job found, release the slot and wait a bit
+                    # Timeout reached, release slot and loop again
                     semaphore.release()
-                    # Prevent busy loop if storage returns None immediately
-                    await sleep(0.1)
+                    # Prevent busy loop if blocking is disabled (e.g. in tests) or failed
+                    if block_time <= 0:
+                        await sleep(0.1)
+
             except CancelledError:
                 break
             except Exception:
                 logger.exception("Error in JobExecutor main loop.")
+                # If an error occurred (e.g. Redis connection lost), sleep briefly to avoid log spam
+                semaphore.release()
                 await sleep(1)
         logger.info("JobExecutor stopped.")
 

avtomatika/health_checker.py

@@ -20,19 +20,37 @@ logger = getLogger(__name__)
 
 
 class HealthChecker:
-    def __init__(self, engine: "OrchestratorEngine"):
+    def __init__(self, engine: "OrchestratorEngine", interval_seconds: int = 600):
+        self.engine = engine
+        self.storage = engine.storage
+        self.interval_seconds = interval_seconds
         self._running = False
+        from uuid import uuid4
+
+        self._instance_id = str(uuid4())
 
     async def run(self):
-        logger.info("HealthChecker is now passive and will not perform active checks.")
+        logger.info(f"HealthChecker started (Active Index Cleanup, Instance ID: {self._instance_id}).")
         self._running = True
         while self._running:
             try:
-                # Sleep for a long time, as this checker is passive.
-                # The loop exists to allow for a clean shutdown.
-                await sleep(3600)
+                # Use distributed lock to ensure only one instance cleans up
+                if await self.storage.acquire_lock(
+                    "global_health_check_lock", self._instance_id, self.interval_seconds - 5
+                ):
+                    try:
+                        await self.storage.cleanup_expired_workers()
+                    finally:
+                        # We don't release the lock immediately to prevent other instances from
+                        # running the same task if the interval is small.
+                        pass
+
+                await sleep(self.interval_seconds)
             except CancelledError:
                 break
+            except Exception:
+                logger.exception("Error in HealthChecker main loop.")
+                await sleep(60)
         logger.info("HealthChecker stopped.")
 
     def stop(self):

avtomatika/history/base.py

@@ -1,25 +1,79 @@
+import asyncio
+import contextlib
 from abc import ABC, abstractmethod
+from logging import getLogger
 from typing import Any
 
+logger = getLogger(__name__)
+
 
 class HistoryStorageBase(ABC):
     """Abstract base class for a history store.
-    Defines the interface for logging job and worker events.
+    Implements buffered asynchronous logging to avoid blocking the main loop.
     """
 
+    def __init__(self):
+        self._queue: asyncio.Queue[tuple[str, dict[str, Any]]] = asyncio.Queue(maxsize=5000)
+        self._worker_task: asyncio.Task | None = None
+
+    async def start(self) -> None:
+        """Starts the background worker for writing logs."""
+        if not self._worker_task:
+            self._worker_task = asyncio.create_task(self._worker())
+            logger.info("HistoryStorage background worker started.")
+
+    async def close(self) -> None:
+        """Stops the background worker and closes resources."""
+        if self._worker_task:
+            self._worker_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._worker_task
+            self._worker_task = None
+            logger.info("HistoryStorage background worker stopped.")
+
     @abstractmethod
-    async def initialize(self):
+    async def initialize(self) -> None:
         """Performs initialization, e.g., creating tables in the DB."""
         raise NotImplementedError
 
+    async def log_job_event(self, event_data: dict[str, Any]) -> None:
+        """Queues a job event for logging."""
+        try:
+            self._queue.put_nowait(("job", event_data))
+        except asyncio.QueueFull:
+            logger.warning("History queue full! Dropping job event.")
+
+    async def log_worker_event(self, event_data: dict[str, Any]) -> None:
+        """Queues a worker event for logging."""
+        try:
+            self._queue.put_nowait(("worker", event_data))
+        except asyncio.QueueFull:
+            logger.warning("History queue full! Dropping worker event.")
+
+    async def _worker(self) -> None:
+        while True:
+            try:
+                kind, data = await self._queue.get()
+                try:
+                    if kind == "job":
+                        await self._persist_job_event(data)
+                    elif kind == "worker":
+                        await self._persist_worker_event(data)
+                except Exception as e:
+                    logger.error(f"Error persisting history event: {e}")
+                finally:
+                    self._queue.task_done()
+            except asyncio.CancelledError:
+                break
+
     @abstractmethod
-    async def log_job_event(self, event_data: dict[str, Any]):
-        """Logs an event related to the job lifecycle."""
+    async def _persist_job_event(self, event_data: dict[str, Any]) -> None:
+        """Actual implementation of writing a job event to storage."""
         raise NotImplementedError
 
     @abstractmethod
-    async def log_worker_event(self, event_data: dict[str, Any]):
-        """Logs an event related to the worker lifecycle."""
+    async def _persist_worker_event(self, event_data: dict[str, Any]) -> None:
+        """Actual implementation of writing a worker event to storage."""
         raise NotImplementedError
 
     @abstractmethod

avtomatika/history/noop.py

@@ -8,20 +8,31 @@ class NoOpHistoryStorage(HistoryStorageBase):
     Used when history storage is not configured.
     """
 
-    async def initialize(self):
-        # Do nothing
+    def __init__(self):
+        super().__init__()
+
+    async def start(self) -> None:
+        pass
+
+    async def close(self) -> None:
+        pass
+
+    async def initialize(self) -> None:
+        pass
+
+    async def log_job_event(self, event_data: dict[str, Any]) -> None:
+        pass
+
+    async def log_worker_event(self, event_data: dict[str, Any]) -> None:
         pass
 
-    async def log_job_event(self, event_data: dict[str, Any]):
-        # Do nothing
+    async def _persist_job_event(self, event_data: dict[str, Any]) -> None:
         pass
 
-    async def log_worker_event(self, event_data: dict[str, Any]):
-        # Do nothing
+    async def _persist_worker_event(self, event_data: dict[str, Any]) -> None:
         pass
 
     async def get_job_history(self, job_id: str) -> list[dict[str, Any]]:
-        # Always return an empty list
         return []
 
     async def get_jobs(self, limit: int = 100, offset: int = 0) -> list[dict[str, Any]]:

avtomatika/history/postgres.py

@@ -46,19 +46,20 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
     """Implementation of the history store based on asyncpg for PostgreSQL."""
 
     def __init__(self, dsn: str, tz_name: str = "UTC"):
+        super().__init__()
         self._dsn = dsn
         self._pool: Pool | None = None
         self.tz_name = tz_name
         self.tz = ZoneInfo(tz_name)
 
-    async def _setup_connection(self, conn: Connection):
+    async def _setup_connection(self, conn: Connection) -> None:
         """Configures the connection session with the correct timezone."""
         try:
             await conn.execute(f"SET TIME ZONE '{self.tz_name}'")
         except PostgresError as e:
             logger.error(f"Failed to set timezone '{self.tz_name}' for PG connection: {e}")
 
-    async def initialize(self):
+    async def initialize(self) -> None:
         """Initializes the connection pool to PostgreSQL and creates tables."""
         try:
             # We use init parameter to configure each new connection in the pool
@@ -75,13 +76,14 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
             logger.error(f"Failed to initialize PostgreSQL history storage: {e}")
             raise
 
-    async def close(self):
-        """Closes the connection pool."""
+    async def close(self) -> None:
+        """Closes the connection pool and background worker."""
+        await super().close()
         if self._pool:
             await self._pool.close()
             logger.info("PostgreSQL history storage connection pool closed.")
 
-    async def log_job_event(self, event_data: dict[str, Any]):
+    async def _persist_job_event(self, event_data: dict[str, Any]) -> None:
         """Logs a job lifecycle event to PostgreSQL."""
         if not self._pool:
             raise RuntimeError("History storage is not initialized.")
@@ -117,7 +119,7 @@ class PostgresHistoryStorage(HistoryStorageBase, ABC):
         except PostgresError as e:
             logger.error(f"Failed to log job event to PostgreSQL: {e}")
 
-    async def log_worker_event(self, event_data: dict[str, Any]):
+    async def _persist_worker_event(self, event_data: dict[str, Any]) -> None:
         """Logs a worker lifecycle event to PostgreSQL."""
         if not self._pool:
             raise RuntimeError("History storage is not initialized.")