avtomatika 1.0b4__py3-none-any.whl → 1.0b6__py3-none-any.whl

avtomatika/scheduler_config_loader.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+from tomllib import load
+from typing import Any
+
+
+@dataclass
+class ScheduledJobConfig:
+    name: str
+    blueprint: str
+    input_data: dict[str, Any]
+    interval_seconds: int | None = None
+    daily_at: str | None = None
+    weekly_days: list[str] | None = None
+    monthly_dates: list[int] | None = None
+    time: str | None = None
+
+
+def load_schedules_from_file(file_path: str) -> list[ScheduledJobConfig]:
+    """Loads scheduled job configurations from a TOML file."""
+    with open(file_path, "rb") as f:
+        data = load(f)
+
+    schedules = []
+    for name, config in data.items():
+        # Skip sections that might be metadata (though TOML structure usually implies all top-level keys are jobs)
+        if not isinstance(config, dict):
+            continue
+
+        schedules.append(
+            ScheduledJobConfig(
+                name=name,
+                blueprint=config.get("blueprint"),
+                input_data=config.get("input_data", {}),
+                interval_seconds=config.get("interval_seconds"),
+                daily_at=config.get("daily_at"),
+                weekly_days=config.get("weekly_days"),
+                monthly_dates=config.get("monthly_dates"),
+                time=config.get("time"),
+            )
+        )
+    return schedules

avtomatika/security.py

@@ -4,11 +4,9 @@ from typing import Any, Awaitable, Callable
 from aiohttp import web
 
 from .config import Config
+from .constants import AUTH_HEADER_CLIENT, AUTH_HEADER_WORKER
 from .storage.base import StorageBackend
 
-AUTH_HEADER_AVTOMATIKA = "X-Avtomatika-Token"
-AUTH_HEADER_WORKER = "X-Worker-Token"
-
 Handler = Callable[[web.Request], Awaitable[web.Response]]
 
 
@@ -21,10 +19,10 @@ def client_auth_middleware_factory(
 
     @web.middleware
     async def middleware(request: web.Request, handler: Handler) -> web.Response:
-        token = request.headers.get(AUTH_HEADER_AVTOMATIKA)
+        token = request.headers.get(AUTH_HEADER_CLIENT)
         if not token:
             return web.json_response(
-                {"error": "Missing X-Avtomatika-Token header"},
+                {"error": f"Missing {AUTH_HEADER_CLIENT} header"},
                 status=401,
             )
 

avtomatika/storage/__init__.py

@@ -1,11 +1,11 @@
-import contextlib
+from contextlib import suppress
 
 from .base import StorageBackend
 from .memory import MemoryStorage
 
 __all__ = ["StorageBackend", "MemoryStorage"]
 
-with contextlib.suppress(ImportError):
+with suppress(ImportError):
     from .redis import RedisStorage  # noqa: F401
 
     __all__.append("RedisStorage")

avtomatika/storage/base.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional
+from typing import Any
 
 
 class StorageBackend(ABC):
@@ -8,7 +8,7 @@ class StorageBackend(ABC):
     """
 
     @abstractmethod
-    async def get_job_state(self, job_id: str) -> Optional[Dict[str, Any]]:
+    async def get_job_state(self, job_id: str) -> dict[str, Any] | None:
         """Get the full state of a job by its ID.
 
         :param job_id: Unique identifier for the job.
@@ -20,8 +20,8 @@ class StorageBackend(ABC):
     async def update_worker_data(
         self,
         worker_id: str,
-        update_data: Dict[str, Any],
-    ) -> Optional[Dict[str, Any]]:
+        update_data: dict[str, Any],
+    ) -> dict[str, Any] | None:
         """Partially update worker information without affecting its TTL.
         Used for background processes like the reputation calculator.
 
@@ -54,9 +54,9 @@ class StorageBackend(ABC):
     async def update_worker_status(
         self,
         worker_id: str,
-        status_update: Dict[str, Any],
+        status_update: dict[str, Any],
         ttl: int,
-    ) -> Optional[Dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Partially update worker information and extend its TTL.
         Used for heartbeat messages.
 
@@ -68,7 +68,7 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def save_job_state(self, job_id: str, state: Dict[str, Any]) -> None:
+    async def save_job_state(self, job_id: str, state: dict[str, Any]) -> None:
         """Save the full state of a job.
 
         :param job_id: Unique identifier for the job.
@@ -80,8 +80,8 @@ class StorageBackend(ABC):
     async def update_job_state(
         self,
         job_id: str,
-        update_data: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        update_data: dict[str, Any],
+    ) -> dict[str, Any]:
         """Partially update the state of a job.
 
         :param job_id: Unique identifier for the job.
@@ -94,7 +94,7 @@ class StorageBackend(ABC):
     async def register_worker(
         self,
         worker_id: str,
-        worker_info: Dict[str, Any],
+        worker_info: dict[str, Any],
         ttl: int,
     ) -> None:
         """Registers a new worker or updates information about an existing one.
@@ -109,7 +109,7 @@ class StorageBackend(ABC):
     async def enqueue_task_for_worker(
         self,
         worker_id: str,
-        task_payload: Dict[str, Any],
+        task_payload: dict[str, Any],
         priority: float,
     ) -> None:
         """Adds a task to the priority queue for a specific worker.
@@ -125,7 +125,7 @@ class StorageBackend(ABC):
         self,
         worker_id: str,
         timeout: int,
-    ) -> Optional[Dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Retrieves the highest priority task from the queue for a worker (blocking operation).
 
         :param worker_id: The ID of the worker for whom to retrieve the task.
@@ -135,7 +135,7 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def get_available_workers(self) -> list[Dict[str, Any]]:
+    async def get_available_workers(self) -> list[dict[str, Any]]:
         """Get a list of all active (not expired) workers.
 
         :return: A list of dictionaries, where each dictionary represents information about a worker.
@@ -165,8 +165,19 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def dequeue_job(self) -> Optional[str]:
-        """Retrieve a job ID from the execution queue (blocking operation)."""
+    async def dequeue_job(self) -> tuple[str, str] | None:
+        """Retrieve a job ID and its message ID from the execution queue.
+
+        :return: A tuple of (job_id, message_id) or None if the timeout has expired.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def ack_job(self, message_id: str) -> None:
+        """Acknowledge successful processing of a job from the queue.
+
+        :param message_id: The identifier of the message to acknowledge.
+        """
         raise NotImplementedError
 
     @abstractmethod
@@ -196,12 +207,12 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def save_client_config(self, token: str, config: Dict[str, Any]) -> None:
+    async def save_client_config(self, token: str, config: dict[str, Any]) -> None:
         """Saves the static configuration of a client."""
         raise NotImplementedError
 
     @abstractmethod
-    async def get_client_config(self, token: str) -> Optional[Dict[str, Any]]:
+    async def get_client_config(self, token: str) -> dict[str, Any] | None:
         """Gets the static configuration of a client."""
         raise NotImplementedError
 
@@ -225,7 +236,7 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def get_priority_queue_stats(self, task_type: str) -> Dict[str, Any]:
+    async def get_priority_queue_stats(self, task_type: str) -> dict[str, Any]:
         """Get statistics on the priority queue for a given task type.
 
         :param task_type: The type of task (used as part of the queue key).
@@ -244,12 +255,12 @@ class StorageBackend(ABC):
         raise NotImplementedError
 
     @abstractmethod
-    async def get_worker_token(self, worker_id: str) -> Optional[str]:
+    async def get_worker_token(self, worker_id: str) -> str | None:
         """Retrieves an individual token for a specific worker."""
         raise NotImplementedError
 
     @abstractmethod
-    async def get_worker_info(self, worker_id: str) -> Optional[Dict[str, Any]]:
+    async def get_worker_info(self, worker_id: str) -> dict[str, Any] | None:
         """Get complete information about a worker by its ID."""
         raise NotImplementedError
 
@@ -258,13 +269,27 @@ class StorageBackend(ABC):
         """Completely clears the storage. Used mainly for tests."""
         raise NotImplementedError
 
-    @abstractmethod
     async def get_active_worker_count(self) -> int:
-        """Get the current number of active (registered) workers.
-        Used for metrics.
+        """Returns the number of currently active workers."""
+        raise NotImplementedError
+
+    async def set_nx_ttl(self, key: str, value: str, ttl: int) -> bool:
+        """
+        Atomically sets key to value if it does not exist.
+        Sets a TTL (in seconds) on the key.
+        Returns True if set, False if already exists.
+        Critical for distributed locking.
         """
         raise NotImplementedError
 
+    async def get_str(self, key: str) -> str | None:
+        """Gets a simple string value from storage."""
+        raise NotImplementedError
+
+    async def set_str(self, key: str, value: str, ttl: int | None = None) -> None:
+        """Sets a simple string value in storage with optional TTL."""
+        raise NotImplementedError
+
     @abstractmethod
     async def acquire_lock(self, key: str, holder_id: str, ttl: int) -> bool:
         """

avtomatika/storage/memory.py

@@ -1,7 +1,7 @@
 from asyncio import Lock, PriorityQueue, Queue, QueueEmpty, wait_for
 from asyncio import TimeoutError as AsyncTimeoutError
 from time import monotonic
-from typing import Any, Dict, List, Optional
+from typing import Any
 
 from .base import StorageBackend
 
@@ -13,35 +13,49 @@ class MemoryStorage(StorageBackend):
     """
 
     def __init__(self):
-        self._jobs: Dict[str, Dict[str, Any]] = {}
-        self._workers: Dict[str, Dict[str, Any]] = {}
-        self._worker_ttls: Dict[str, float] = {}
-        self._worker_task_queues: Dict[str, PriorityQueue] = {}
+        self._jobs: dict[str, dict[str, Any]] = {}
+        self._workers: dict[str, dict[str, Any]] = {}
+        self._worker_ttls: dict[str, float] = {}
+        self._worker_task_queues: dict[str, PriorityQueue] = {}
         self._job_queue = Queue()
-        self._quarantine_queue: List[str] = []
-        self._watched_jobs: Dict[str, float] = {}
-        self._client_configs: Dict[str, Dict[str, Any]] = {}
-        self._quotas: Dict[str, int] = {}
-        self._worker_tokens: Dict[str, str] = {}
-        self._generic_keys: Dict[str, Any] = {}
-        self._generic_key_ttls: Dict[str, float] = {}
-        self._locks: Dict[str, tuple[str, float]] = {}  # key -> (holder_id, expiry_time)
+        self._quarantine_queue: list[str] = []
+        self._watched_jobs: dict[str, float] = {}
+        self._client_configs: dict[str, dict[str, Any]] = {}
+        self._quotas: dict[str, int] = {}
+        self._worker_tokens: dict[str, str] = {}
+        self._generic_keys: dict[str, Any] = {}
+        self._generic_key_ttls: dict[str, float] = {}
+        self._locks: dict[str, tuple[str, float]] = {}
 
         self._lock = Lock()
 
-    async def get_job_state(self, job_id: str) -> Optional[Dict[str, Any]]:
+    async def get_job_state(self, job_id: str) -> dict[str, Any] | None:
         async with self._lock:
             return self._jobs.get(job_id)
 
-    async def save_job_state(self, job_id: str, state: Dict[str, Any]) -> None:
+    async def _clean_expired(self):
+        """Helper to remove expired keys."""
+        now = monotonic()
+
+        expired_generic = [k for k, t in self._generic_key_ttls.items() if t < now]
+        for k in expired_generic:
+            self._generic_key_ttls.pop(k, None)
+            self._generic_keys.pop(k, None)
+
+        expired_workers = [k for k, t in self._worker_ttls.items() if t < now]
+        for k in expired_workers:
+            self._worker_ttls.pop(k, None)
+            self._workers.pop(k, None)
+
+    async def save_job_state(self, job_id: str, state: dict[str, Any]):
         async with self._lock:
             self._jobs[job_id] = state
 
     async def update_job_state(
         self,
         job_id: str,
-        update_data: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        update_data: dict[str, Any],
+    ) -> dict[str, Any]:
         async with self._lock:
             if job_id not in self._jobs:
                 self._jobs[job_id] = {}
@@ -51,7 +65,7 @@ class MemoryStorage(StorageBackend):
     async def register_worker(
         self,
         worker_id: str,
-        worker_info: Dict[str, Any],
+        worker_info: dict[str, Any],
         ttl: int,
     ) -> None:
         """Registers a worker and creates a task queue for it."""
@@ -66,21 +80,20 @@ class MemoryStorage(StorageBackend):
     async def enqueue_task_for_worker(
         self,
         worker_id: str,
-        task_payload: Dict[str, Any],
+        task_payload: dict[str, Any],
         priority: float,
     ) -> None:
         """Puts a task on the priority queue for a worker."""
         async with self._lock:
             if worker_id not in self._worker_task_queues:
                 self._worker_task_queues[worker_id] = PriorityQueue()
-        # asyncio.PriorityQueue is a min-heap, so we invert the priority
         await self._worker_task_queues[worker_id].put((-priority, task_payload))
 
     async def dequeue_task_for_worker(
         self,
         worker_id: str,
         timeout: int,
-    ) -> Optional[Dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         """Retrieves a task from the worker's priority queue with a timeout."""
         queue = None
         async with self._lock:
@@ -104,9 +117,9 @@ class MemoryStorage(StorageBackend):
     async def update_worker_status(
         self,
         worker_id: str,
-        status_update: Dict[str, Any],
+        status_update: dict[str, Any],
         ttl: int,
-    ) -> Optional[Dict[str, Any]]:
+    ) -> dict[str, Any] | None:
         async with self._lock:
             if worker_id in self._workers:
                 self._workers[worker_id].update(status_update)
@@ -117,8 +130,8 @@ class MemoryStorage(StorageBackend):
     async def update_worker_data(
         self,
         worker_id: str,
-        update_data: Dict[str, Any],
-    ) -> Optional[Dict[str, Any]]:
+        update_data: dict[str, Any],
+    ) -> dict[str, Any] | None:
         async with self._lock:
             if worker_id in self._workers:
                 self._workers[worker_id].update(update_data)
@@ -155,13 +168,17 @@ class MemoryStorage(StorageBackend):
     async def enqueue_job(self, job_id: str) -> None:
         await self._job_queue.put(job_id)
 
-    async def dequeue_job(self) -> str | None:
+    async def dequeue_job(self) -> tuple[str, str] | None:
         """Waits indefinitely for a job ID from the queue and returns it.
-        This simulates the blocking behavior of Redis's BLPOP.
+        Returns a tuple of (job_id, message_id). In MemoryStorage, message_id is dummy.
         """
         job_id = await self._job_queue.get()
         self._job_queue.task_done()
-        return job_id
+        return job_id, "memory-msg-id"
+
+    async def ack_job(self, message_id: str) -> None:
+        """No-op for MemoryStorage as it doesn't support persistent streams."""
+        pass
 
     async def quarantine_job(self, job_id: str) -> None:
         async with self._lock:
@@ -187,11 +204,11 @@ class MemoryStorage(StorageBackend):
             self._generic_key_ttls[key] = now + ttl
             return self._generic_keys[key]
 
-    async def save_client_config(self, token: str, config: Dict[str, Any]) -> None:
+    async def save_client_config(self, token: str, config: dict[str, Any]) -> None:
         async with self._lock:
             self._client_configs[token] = config
 
-    async def get_client_config(self, token: str) -> Optional[Dict[str, Any]]:
+    async def get_client_config(self, token: str) -> dict[str, Any] | None:
         async with self._lock:
             return self._client_configs.get(token)
 
@@ -217,7 +234,6 @@ class MemoryStorage(StorageBackend):
             self._workers.clear()
             self._worker_ttls.clear()
             self._worker_task_queues.clear()
-            # Empty the queue
             while not self._job_queue.empty():
                 try:
                     self._job_queue.get_nowait()
@@ -232,17 +248,38 @@ class MemoryStorage(StorageBackend):
             self._locks.clear()
 
     async def get_job_queue_length(self) -> int:
-        # No lock needed for asyncio.Queue.qsize()
         return self._job_queue.qsize()
 
     async def get_active_worker_count(self) -> int:
         async with self._lock:
-            now = monotonic()
-            # Create a copy of keys to avoid issues with concurrent modifications
-            worker_ids = list(self._workers.keys())
-            return sum(self._worker_ttls.get(worker_id, 0) > now for worker_id in worker_ids)
+            await self._clean_expired()
+            return len(self._workers)
+
+    async def set_nx_ttl(self, key: str, value: str, ttl: int) -> bool:
+        async with self._lock:
+            await self._clean_expired()
+            if key in self._generic_keys:
+                return False
 
-    async def get_worker_info(self, worker_id: str) -> Optional[Dict[str, Any]]:
+            self._generic_keys[key] = value
+            self._generic_key_ttls[key] = monotonic() + ttl
+            return True
+
+    async def get_str(self, key: str) -> str | None:
+        async with self._lock:
+            await self._clean_expired()
+            val = self._generic_keys.get(key)
+            return str(val) if val is not None else None
+
+    async def set_str(self, key: str, value: str, ttl: int | None = None) -> None:
+        async with self._lock:
+            self._generic_keys[key] = value
+            if ttl:
+                self._generic_key_ttls[key] = monotonic() + ttl
+            else:
+                self._generic_key_ttls.pop(key, None)
+
+    async def get_worker_info(self, worker_id: str) -> dict[str, Any] | None:
         async with self._lock:
             return self._workers.get(worker_id)
 
@@ -250,7 +287,7 @@ class MemoryStorage(StorageBackend):
         async with self._lock:
             self._worker_tokens[worker_id] = token
 
-    async def get_worker_token(self, worker_id: str) -> Optional[str]:
+    async def get_worker_token(self, worker_id: str) -> str | None:
         async with self._lock:
             return self._worker_tokens.get(worker_id)
 
@@ -258,7 +295,7 @@ class MemoryStorage(StorageBackend):
         key = f"task_cancel:{task_id}"
         await self.increment_key_with_ttl(key, 3600)
 
-    async def get_priority_queue_stats(self, task_type: str) -> Dict[str, Any]:
+    async def get_priority_queue_stats(self, task_type: str) -> dict[str, Any]:
         """
         Returns empty data, as `asyncio.PriorityQueue` does not
         support introspection to get statistics.
@@ -278,14 +315,8 @@ class MemoryStorage(StorageBackend):
         async with self._lock:
             now = monotonic()
             current_lock = self._locks.get(key)
-
-            # If lock exists and hasn't expired
             if current_lock and current_lock[1] > now:
-                # If explicitly owned by us, we can extend/re-enter (optional behavior)
-                # But for strict locking, if it's held, return False (unless it's us? let's simpler: just False if held)
                 return False
-
-            # Acquire lock
             self._locks[key] = (holder_id, now + ttl)
             return True
 
@@ -294,7 +325,6 @@ class MemoryStorage(StorageBackend):
             current_lock = self._locks.get(key)
             if current_lock:
                 owner, expiry = current_lock
-                # Only release if we are the owner
                 if owner == holder_id:
                     del self._locks[key]
                     return True