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

avtomatika/history/sqlite.py

@@ -49,11 +49,12 @@ class SQLiteHistoryStorage(HistoryStorageBase):
     """
 
     def __init__(self, db_path: str, tz_name: str = "UTC"):
+        super().__init__()
         self._db_path = db_path
         self._conn: Connection | None = None
         self.tz = ZoneInfo(tz_name)
 
-    async def initialize(self):
+    async def initialize(self) -> None:
         """Initializes the database connection and creates tables if they don't exist."""
         try:
             self._conn = await connect(self._db_path)
@@ -68,8 +69,9 @@ class SQLiteHistoryStorage(HistoryStorageBase):
             logger.error(f"Failed to initialize SQLite history storage: {e}")
             raise
 
-    async def close(self):
-        """Closes the database connection."""
+    async def close(self) -> None:
+        """Closes the database connection and background worker."""
+        await super().close()
         if self._conn:
             await self._conn.close()
             logger.info("SQLite history storage connection closed.")
@@ -91,7 +93,7 @@ class SQLiteHistoryStorage(HistoryStorageBase):
 
         return item
 
-    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 the job_history table."""
         if not self._conn:
             raise RuntimeError("History storage is not initialized.")
@@ -128,7 +130,7 @@ class SQLiteHistoryStorage(HistoryStorageBase):
         except Error as e:
             logger.error(f"Failed to log job event: {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 the worker_history table."""
         if not self._conn:
             raise RuntimeError("History storage is not initialized.")

avtomatika/metrics.py

@@ -12,7 +12,7 @@ task_queue_length: Gauge
 active_workers: Gauge
 
 
-def init_metrics():
+def init_metrics() -> None:
     """
     Initializes Prometheus metrics.
     Uses a registry check for idempotency, which is important for tests.

avtomatika/reputation.py

@@ -52,48 +52,54 @@ class ReputationCalculator:
     async def calculate_all_reputations(self):
         """Calculates and updates the reputation for all active workers."""
         logger.info("Starting reputation calculation for all workers...")
-        workers = await self.storage.get_available_workers()
-        if not workers:
+
+        # Get only IDs of active workers to avoid O(N) scan of all data
+        worker_ids = await self.storage.get_active_worker_ids()
+
+        if not worker_ids:
             logger.info("No active workers found for reputation calculation.")
             return
 
-        for worker in workers:
-            worker_id = worker.get("worker_id")
-            if not worker_id:
-                continue
-
-            history = await self.history_storage.get_worker_history(
-                worker_id,
-                since_days=REPUTATION_HISTORY_DAYS,
-            )
-
-            # Count only task completion events
-            task_finished_events = [event for event in history if event.get("event_type") == "task_finished"]
-
-            if not task_finished_events:
-                # If there is no history, the reputation does not change (remains 1.0 by default)
-                continue
-
-            successful_tasks = 0
-            for event in task_finished_events:
-                # Extract the result from the snapshot
-                snapshot = event.get("context_snapshot", {})
-                result = snapshot.get("result", {})
-                if result.get("status") == "success":
-                    successful_tasks += 1
-
-            total_tasks = len(task_finished_events)
-            new_reputation = successful_tasks / total_tasks if total_tasks > 0 else 1.0
-
-            # Round for cleanliness
-            new_reputation = round(new_reputation, 4)
-
-            logger.info(
-                f"Updating reputation for worker {worker_id}: {worker.get('reputation')} -> {new_reputation}",
-            )
-            await self.storage.update_worker_data(
-                worker_id,
-                {"reputation": new_reputation},
-            )
+        logger.info(f"Recalculating reputation for {len(worker_ids)} workers.")
+
+        for worker_id in worker_ids:
+            if not self._running:
+                break
+
+            try:
+                history = await self.history_storage.get_worker_history(
+                    worker_id,
+                    since_days=REPUTATION_HISTORY_DAYS,
+                )
+
+                # Count only task completion events
+                task_finished_events = [event for event in history if event.get("event_type") == "task_finished"]
+
+                if not task_finished_events:
+                    # If there is no history, skip to next worker
+                    continue
+
+                successful_tasks = 0
+                for event in task_finished_events:
+                    # Extract the result from the snapshot
+                    snapshot = event.get("context_snapshot", {})
+                    result = snapshot.get("result", {})
+                    if result.get("status") == "success":
+                        successful_tasks += 1
+
+                total_tasks = len(task_finished_events)
+                new_reputation = successful_tasks / total_tasks if total_tasks > 0 else 1.0
+                new_reputation = round(new_reputation, 4)
+
+                await self.storage.update_worker_data(
+                    worker_id,
+                    {"reputation": new_reputation},
+                )
+
+                # Throttling: Small sleep to prevent DB spikes
+                await sleep(0.1)
+
+            except Exception as e:
+                logger.error(f"Failed to calculate reputation for worker {worker_id}: {e}")
 
         logger.info("Reputation calculation finished.")

avtomatika/s3.py

@@ -0,0 +1,379 @@
+from asyncio import Semaphore, gather, to_thread
+from logging import getLogger
+from os import sep, walk
+from pathlib import Path
+from shutil import rmtree
+from typing import Any
+
+from aiofiles import open as aiopen
+from obstore import delete_async, get_async, put_async
+from obstore import list as obstore_list
+from obstore.store import S3Store
+from orjson import dumps, loads
+from rxon.blob import calculate_config_hash, parse_uri
+from rxon.exceptions import IntegrityError
+
+from .config import Config
+from .history.base import HistoryStorageBase
+
+logger = getLogger(__name__)
+
+try:
+    HAS_S3_LIBS = True
+except ImportError:
+    HAS_S3_LIBS = False
+    S3Store = Any
+
+
+class TaskFiles:
+    """
+    Manages files for a specific job, ensuring full compatibility with avtomatika-worker.
+    Supports recursive directory download/upload and non-blocking I/O.
+    """
+
+    def __init__(
+        self,
+        store: "S3Store",
+        bucket: str,
+        job_id: str,
+        base_local_dir: str | Path,
+        semaphore: Semaphore,
+        history: HistoryStorageBase | None = None,
+    ):
+        self._store = store
+        self._bucket = bucket
+        self._job_id = job_id
+        self._history = history
+        self._s3_prefix = f"jobs/{job_id}/"
+        self.local_dir = Path(base_local_dir) / job_id
+        self._semaphore = semaphore
+
+    def _ensure_local_dir(self) -> None:
+        if not self.local_dir.exists():
+            self.local_dir.mkdir(parents=True, exist_ok=True)
+
+    def path(self, filename: str) -> Path:
+        """Returns local path for a filename, ensuring the directory exists."""
+        self._ensure_local_dir()
+        clean_name = filename.split("/")[-1] if "://" in filename else filename.lstrip("/")
+        return self.local_dir / clean_name
+
+    async def _download_single_file(
+        self,
+        key: str,
+        local_path: Path,
+        expected_size: int | None = None,
+        expected_hash: str | None = None,
+    ) -> dict[str, Any]:
+        """Downloads a single file safely using semaphore and streaming.
+        Returns metadata (size, etag).
+        """
+        if not local_path.parent.exists():
+            await to_thread(local_path.parent.mkdir, parents=True, exist_ok=True)
+
+        async with self._semaphore:
+            response = await get_async(self._store, key)
+            meta = response.meta
+            file_size = meta.size
+            etag = meta.e_tag.strip('"') if meta.e_tag else None
+
+            if expected_size is not None and file_size != expected_size:
+                raise IntegrityError(f"File size mismatch for {key}: expected {expected_size}, got {file_size}")
+
+            if expected_hash is not None and etag and expected_hash != etag:
+                raise IntegrityError(f"Integrity mismatch for {key}: expected ETag {expected_hash}, got {etag}")
+
+            stream = response.stream()
+            async with aiopen(local_path, "wb") as f:
+                async for chunk in stream:
+                    await f.write(chunk)
+
+            return {"size": file_size, "etag": etag}
+
+    async def download(
+        self,
+        name_or_uri: str,
+        local_name: str | None = None,
+        verify_meta: dict[str, Any] | None = None,
+    ) -> Path:
+        """
+        Downloads a file or directory (recursively).
+        If URI ends with '/', it treats it as a directory.
+        """
+        bucket, key, is_dir = parse_uri(name_or_uri, self._bucket, self._s3_prefix)
+        verify_meta = verify_meta or {}
+
+        if local_name:
+            target_path = self.path(local_name)
+        else:
+            suffix = key.replace(self._s3_prefix, "", 1) if key.startswith(self._s3_prefix) else key.split("/")[-1]
+            target_path = self.local_dir / suffix
+
+        if is_dir:
+            logger.info(f"Recursive download: s3://{bucket}/{key} -> {target_path}")
+            entries = await to_thread(lambda: list(obstore_list(self._store, prefix=key)))
+
+            tasks = []
+            for entry in entries:
+                s3_key = entry["path"]
+                rel_path = s3_key[len(key) :]
+                if not rel_path:
+                    continue
+
+                local_file_path = target_path / rel_path
+                tasks.append(self._download_single_file(s3_key, local_file_path))
+
+            if tasks:
+                results = await gather(*tasks)
+                total_size = sum(r["size"] for r in results)
+                await self._log_event(
+                    "download_dir",
+                    f"s3://{bucket}/{key}",
+                    str(target_path),
+                    metadata={"total_size": total_size, "file_count": len(results)},
+                )
+            else:
+                await self._log_event(
+                    "download_dir",
+                    f"s3://{bucket}/{key}",
+                    str(target_path),
+                    metadata={"total_size": 0, "file_count": 0},
+                )
+            return target_path
+        else:
+            logger.debug(f"Downloading s3://{bucket}/{key} -> {target_path}")
+            meta = await self._download_single_file(
+                key,
+                target_path,
+                expected_size=verify_meta.get("size"),
+                expected_hash=verify_meta.get("hash"),
+            )
+            await self._log_event("download", f"s3://{bucket}/{key}", str(target_path), metadata=meta)
+            return target_path
+
+    async def _upload_single_file(self, local_path: Path, s3_key: str) -> dict[str, Any]:
+        """Uploads a single file safely using semaphore. Returns S3 metadata."""
+        async with self._semaphore:
+            file_size = local_path.stat().st_size
+            async with aiopen(local_path, "rb") as f:
+                content = await f.read()
+            result = await put_async(self._store, s3_key, content)
+            etag = result.e_tag.strip('"') if result.e_tag else None
+            return {"size": file_size, "etag": etag}
+
+    async def upload(self, local_name: str, remote_name: str | None = None) -> str:
+        """
+        Uploads a file or directory recursively.
+        If local_name points to a directory, it uploads all contents.
+        """
+        local_path = self.path(local_name)
+
+        if local_path.is_dir():
+            base_remote = (remote_name or local_name).lstrip("/")
+            if not base_remote.endswith("/"):
+                base_remote += "/"
+
+            target_prefix = f"{self._s3_prefix}{base_remote}"
+            logger.info(f"Recursive upload: {local_path} -> s3://{self._bucket}/{target_prefix}")
+
+            def collect_files():
+                files_to_upload = []
+                for root, _, files in walk(local_path):
+                    for file in files:
+                        abs_path = Path(root) / file
+                        rel_path = abs_path.relative_to(local_path)
+                        s3_key = f"{target_prefix}{str(rel_path).replace(sep, '/')}"
+                        files_to_upload.append((abs_path, s3_key))
+                return files_to_upload
+
+            files_map = await to_thread(collect_files)
+
+            tasks = [self._upload_single_file(lp, k) for lp, k in files_map]
+            if tasks:
+                results = await gather(*tasks)
+                total_size = sum(r["size"] for r in results)
+                metadata = {"total_size": total_size, "file_count": len(results)}
+            else:
+                metadata = {"total_size": 0, "file_count": 0}
+
+            uri = f"s3://{self._bucket}/{target_prefix}"
+            await self._log_event("upload_dir", uri, str(local_path), metadata=metadata)
+            return uri
+
+        elif local_path.exists():
+            target_key = f"{self._s3_prefix}{(remote_name or local_name).lstrip('/')}"
+            logger.debug(f"Uploading {local_path} -> s3://{self._bucket}/{target_key}")
+
+            meta = await self._upload_single_file(local_path, target_key)
+
+            uri = f"s3://{self._bucket}/{target_key}"
+            await self._log_event("upload", uri, str(local_path), metadata=meta)
+            return uri
+        else:
+            raise FileNotFoundError(f"Local file/dir not found: {local_path}")
+
+    async def read_text(self, name_or_uri: str) -> str:
+        bucket, key, _ = parse_uri(name_or_uri, self._bucket, self._s3_prefix)
+        filename = key.split("/")[-1]
+        local_path = self.path(filename)
+
+        if not local_path.exists():
+            await self.download(name_or_uri)
+
+        async with aiopen(local_path, "r", encoding="utf-8") as f:
+            return await f.read()
+
+    async def read_json(self, name_or_uri: str) -> Any:
+        bucket, key, _ = parse_uri(name_or_uri, self._bucket, self._s3_prefix)
+        filename = key.split("/")[-1]
+        local_path = self.path(filename)
+
+        if not local_path.exists():
+            await self.download(name_or_uri)
+
+        async with aiopen(local_path, "rb") as f:
+            content = await f.read()
+            return loads(content)
+
+    async def write_json(self, filename: str, data: Any, upload: bool = True) -> str:
+        """Writes JSON locally (binary mode) and optionally uploads to S3."""
+        local_path = self.path(filename)
+        json_bytes = dumps(data)
+
+        async with aiopen(local_path, "wb") as f:
+            await f.write(json_bytes)
+
+        if upload:
+            return await self.upload(filename)
+        return f"file://{local_path}"
+
+    async def write_text(self, filename: str, text: str, upload: bool = True) -> Path:
+        local_path = self.path(filename)
+        async with aiopen(local_path, "w", encoding="utf-8") as f:
+            await f.write(text)
+
+        if upload:
+            await self.upload(filename)
+
+        return local_path
+
+    async def cleanup(self) -> None:
+        """Full cleanup of S3 prefix and local job directory."""
+        logger.info(f"Cleanup for job {self._job_id}...")
+        try:
+            entries = await to_thread(lambda: list(obstore_list(self._store, prefix=self._s3_prefix)))
+            paths_to_delete = [entry["path"] for entry in entries]
+            if paths_to_delete:
+                await delete_async(self._store, paths_to_delete)
+        except Exception as e:
+            logger.error(f"S3 cleanup error: {e}")
+
+        if self.local_dir.exists():
+            await to_thread(rmtree, self.local_dir)
+
+    async def _log_event(
+        self,
+        operation: str,
+        file_uri: str,
+        local_path: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
+        if not self._history:
+            return
+
+        try:
+            context_snapshot = {
+                "operation": operation,
+                "s3_uri": file_uri,
+                "local_path": str(local_path),
+            }
+            if metadata:
+                context_snapshot.update(metadata)
+
+            await self._history.log_job_event(
+                {
+                    "job_id": self._job_id,
+                    "event_type": "s3_operation",
+                    "state": "running",
+                    "context_snapshot": context_snapshot,
+                }
+            )
+        except Exception as e:
+            logger.warning(f"Failed to log S3 event: {e}")
+
+
+class S3Service:
+    """
+    Central service for S3 operations.
+    Initializes the Store and provides TaskFiles instances.
+    """
+
+    def __init__(self, config: Config, history: HistoryStorageBase | None = None):
+        self.config = config
+        self._history = history
+        self._store: S3Store | None = None
+        self._semaphore: Semaphore | None = None
+
+        self._config_present = bool(config.S3_ENDPOINT_URL and config.S3_ACCESS_KEY and config.S3_SECRET_KEY)
+
+        if self._config_present:
+            if HAS_S3_LIBS:
+                self._enabled = True
+                self._initialize_store()
+            else:
+                logger.error(
+                    "S3 configuration found, but 'avtomatika[s3]' extra dependencies are not installed. "
+                    "S3 support will be disabled. Install with: pip install 'avtomatika[s3]'"
+                )
+                self._enabled = False
+        else:
+            self._enabled = False
+            if any([config.S3_ENDPOINT_URL, config.S3_ACCESS_KEY, config.S3_SECRET_KEY]):
+                logger.warning("Partial S3 configuration found. S3 support disabled.")
+
+    def _initialize_store(self) -> None:
+        try:
+            self._store = S3Store(
+                bucket=self.config.S3_DEFAULT_BUCKET,
+                access_key_id=self.config.S3_ACCESS_KEY,
+                secret_access_key=self.config.S3_SECRET_KEY,
+                region=self.config.S3_REGION,
+                endpoint=self.config.S3_ENDPOINT_URL,
+                allow_http="http://" in self.config.S3_ENDPOINT_URL,
+                force_path_style=True,
+            )
+            self._semaphore = Semaphore(self.config.S3_MAX_CONCURRENCY)
+            logger.info(
+                f"S3Service initialized (Endpoint: {self.config.S3_ENDPOINT_URL}, "
+                f"Bucket: {self.config.S3_DEFAULT_BUCKET}, "
+                f"Max Concurrency: {self.config.S3_MAX_CONCURRENCY})"
+            )
+        except Exception as e:
+            logger.error(f"Failed to initialize S3 Store: {e}")
+            self._enabled = False
+
+    def get_config_hash(self) -> str | None:
+        """Returns a hash of the current S3 configuration for consistency checks."""
+        if not self._enabled:
+            return None
+        return calculate_config_hash(
+            self.config.S3_ENDPOINT_URL,
+            self.config.S3_ACCESS_KEY,
+            self.config.S3_DEFAULT_BUCKET,
+        )
+
+    def get_task_files(self, job_id: str) -> TaskFiles | None:
+        if not self._enabled or not self._store or not self._semaphore:
+            return None
+
+        return TaskFiles(
+            self._store,
+            self.config.S3_DEFAULT_BUCKET,
+            job_id,
+            self.config.TASK_FILES_DIR,
+            self._semaphore,
+            self._history,
+        )
+
+    async def close(self) -> None:
+        pass

avtomatika/security.py

@@ -10,6 +10,62 @@ from .storage.base import StorageBackend
 Handler = Callable[[web.Request], Awaitable[web.Response]]
 
 
+async def verify_worker_auth(
+    storage: StorageBackend,
+    config: Config,
+    token: str | None,
+    cert_identity: str | None,
+    worker_id_hint: str | None,
+) -> str:
+    """
+    Verifies worker authentication using token or mTLS.
+    Returns authenticated worker_id.
+    Raises ValueError (400), PermissionError (401/403) on failure.
+    """
+    # mTLS Check
+    if cert_identity:
+        if worker_id_hint and cert_identity != worker_id_hint:
+            raise PermissionError(
+                f"Unauthorized: Certificate CN '{cert_identity}' does not match worker_id '{worker_id_hint}'"
+            )
+        return cert_identity
+
+    # Token Check
+    if not token:
+        raise PermissionError(f"Missing {AUTH_HEADER_WORKER} header or client certificate")
+
+    hashed_provided_token = sha256(token.encode()).hexdigest()
+
+    # STS Access Token
+    token_worker_id = await storage.verify_worker_access_token(hashed_provided_token)
+    if token_worker_id:
+        if worker_id_hint and token_worker_id != worker_id_hint:
+            raise PermissionError(
+                f"Unauthorized: Access Token belongs to '{token_worker_id}', but request is for '{worker_id_hint}'"
+            )
+        return token_worker_id
+
+    # Individual/Global Token
+    if not worker_id_hint:
+        if config.GLOBAL_WORKER_TOKEN and token == config.GLOBAL_WORKER_TOKEN:
+            return "unknown_authenticated_by_global_token"
+
+        raise PermissionError("Unauthorized: Invalid token or missing worker_id hint")
+
+    # Individual Token for specific worker
+    expected_token_hash = await storage.get_worker_token(worker_id_hint)
+    if expected_token_hash:
+        if hashed_provided_token == expected_token_hash:
+            return worker_id_hint
+        raise PermissionError("Unauthorized: Invalid individual worker token")
+
+    # Global Token Fallback
+    if config.GLOBAL_WORKER_TOKEN and token == config.GLOBAL_WORKER_TOKEN:
+        return worker_id_hint
+
+    raise PermissionError("Unauthorized: No valid token found")
+
+
 def client_auth_middleware_factory(
     storage: StorageBackend,
 ) -> Any:
@@ -38,77 +94,3 @@ def client_auth_middleware_factory(
         return await handler(request)
 
     return middleware
-
-
-def worker_auth_middleware_factory(
-    storage: StorageBackend,
-    config: Config,
-) -> Any:
-    """
-    Middleware factory for worker authentication.
-    It supports both individual tokens and a global fallback token for backward compatibility.
-    It also attaches the authenticated worker_id to the request.
-    """
-
-    @web.middleware
-    async def middleware(request: web.Request, handler: Handler) -> web.Response:
-        provided_token = request.headers.get(AUTH_HEADER_WORKER)
-        if not provided_token:
-            return web.json_response(
-                {"error": f"Missing {AUTH_HEADER_WORKER} header"},
-                status=401,
-            )
-
-        worker_id = request.match_info.get("worker_id")
-        data = None
-
-        # For specific endpoints, worker_id is in the body.
-        # We need to read the body here, which can be tricky as it's a stream.
-        # We clone the request to allow the handler to read the body again.
-        if not worker_id and (request.path.endswith("/register") or request.path.endswith("/tasks/result")):
-            try:
-                cloned_request = request.clone()
-                data = await cloned_request.json()
-                worker_id = data.get("worker_id")
-                # Attach the parsed data to the request so the handler doesn't need to re-parse
-                if request.path.endswith("/register"):
-                    request["worker_registration_data"] = data
-            except Exception:
-                return web.json_response({"error": "Invalid JSON body"}, status=400)
-
-        # If no worker_id could be determined from path or body, we can only validate against the global token.
-        if not worker_id:
-            if provided_token == config.GLOBAL_WORKER_TOKEN:
-                # We don't know the worker_id, so we can't attach it.
-                return await handler(request)
-            else:
-                return web.json_response(
-                    {"error": "Unauthorized: Invalid token or missing worker_id"},
-                    status=401,
-                )
-
-        # --- Individual Token Check ---
-        expected_token_hash = await storage.get_worker_token(worker_id)
-        if expected_token_hash:
-            hashed_provided_token = sha256(provided_token.encode()).hexdigest()
-            if hashed_provided_token == expected_token_hash:
-                request["worker_id"] = worker_id  # Attach authenticated worker_id
-                return await handler(request)
-            else:
-                # If an individual token exists, we do not fall back to the global token.
-                return web.json_response(
-                    {"error": "Unauthorized: Invalid individual worker token"},
-                    status=401,
-                )
-
-        # --- Global Token Fallback ---
-        if config.GLOBAL_WORKER_TOKEN and provided_token == config.GLOBAL_WORKER_TOKEN:
-            request["worker_id"] = worker_id  # Attach authenticated worker_id
-            return await handler(request)
-
-        return web.json_response(
-            {"error": "Unauthorized: No valid token found"},
-            status=401,
-        )
-
-    return middleware