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

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
 

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,35 @@ 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 save_job_state(self, job_id: str, state: dict[str, Any]) -> None:
         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 +51,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 +66,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 +103,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 +116,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 +154,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 +190,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 +220,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 +234,15 @@ 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)
 
-    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:
         async with self._lock:
             return self._workers.get(worker_id)
 
@@ -250,7 +250,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 +258,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 +278,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 +288,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