coati-payroll 0.0.2__py3-none-any.whl

coati_payroll/queue/drivers/dramatiq_driver.py

@@ -0,0 +1,268 @@
+# Copyright 2025 BMO Soluciones, S.A.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dramatiq driver for high-performance queue processing with Redis."""
+
+from __future__ import annotations
+
+from typing import Any, Callable
+
+from coati_payroll.log import log
+from coati_payroll.queue.driver import QueueDriver
+
+
+class DramatiqDriver(QueueDriver):
+    """Queue driver using Dramatiq with Redis backend.
+
+    This driver provides high-performance, distributed job processing
+    suitable for production environments with Redis available.
+
+    Features:
+    - Multi-threaded workers
+    - Automatic retries with exponential backoff
+    - Distributed processing across multiple workers
+    - Results backend (optional)
+    """
+
+    def __init__(self, redis_url: str | None = None):
+        """Initialize Dramatiq driver.
+
+        Args:
+            redis_url: Redis connection URL (default: redis://localhost:6379/0)
+        """
+        self._redis_url = redis_url or "redis://localhost:6379/0"
+        self._broker = None
+        self._tasks = {}
+        self._available = self._initialize_broker()
+
+    def _initialize_broker(self) -> bool:
+        """Initialize the Dramatiq broker.
+
+        Returns:
+            True if broker initialized successfully, False otherwise
+        """
+        try:
+            import dramatiq
+            from dramatiq.brokers.redis import RedisBroker
+            from dramatiq.middleware import (
+                AgeLimit,
+                Retries,
+                TimeLimit,
+            )
+
+            # Test Redis connection
+            import redis
+
+            client = redis.from_url(self._redis_url, socket_connect_timeout=2)
+            client.ping()
+
+            # Configure broker with middleware
+            self._broker = RedisBroker(url=self._redis_url)
+
+            # Add middleware for retries, age limits, and time limits
+            self._broker.add_middleware(Retries(max_retries=3, min_backoff=15000, max_backoff=86400000))
+            self._broker.add_middleware(TimeLimit(time_limit=3600000))  # 1 hour max
+            self._broker.add_middleware(AgeLimit(max_age=86400000))  # 24 hours max age
+
+            # Set as default broker
+            dramatiq.set_broker(self._broker)
+
+            log.info(f"Dramatiq driver initialized with Redis at {self._redis_url}")
+            return True
+
+        except ImportError as e:
+            log.warning(f"Dramatiq not available: {e}")
+            return False
+        except Exception as e:
+            log.warning(f"Failed to connect to Redis for Dramatiq: {e}")
+            return False
+
+    def enqueue(self, task_name: str, *args: Any, delay: int | None = None, **kwargs: Any) -> Any:
+        """Enqueue a task for background processing.
+
+        Args:
+            task_name: Name of the registered task
+            *args: Positional arguments for the task
+            delay: Optional delay in seconds before execution
+            **kwargs: Keyword arguments for the task
+
+        Returns:
+            Dramatiq message object
+
+        Raises:
+            ValueError: If task is not registered
+            RuntimeError: If driver is not available
+        """
+        if not self._available:
+            raise RuntimeError("Dramatiq driver is not available")
+
+        if task_name not in self._tasks:
+            raise ValueError(f"Task '{task_name}' not registered")
+
+        task = self._tasks[task_name]
+
+        if delay:
+            # Convert seconds to milliseconds for Dramatiq
+            return task.send_with_options(args=args, kwargs=kwargs, delay=delay * 1000)
+        else:
+            return task.send(*args, **kwargs)
+
+    def register_task(
+        self,
+        func: Callable,
+        name: str | None = None,
+        max_retries: int = 3,
+        min_backoff: int = 15000,
+        max_backoff: int = 86400000,
+    ) -> Callable:
+        """Register a function as a Dramatiq actor.
+
+        Args:
+            func: Function to register
+            name: Optional task name (defaults to function name)
+            max_retries: Maximum retry attempts
+            min_backoff: Minimum backoff in milliseconds
+            max_backoff: Maximum backoff in milliseconds
+
+        Returns:
+            Dramatiq actor that can be called or enqueued
+        """
+        if not self._available:
+            log.warning(f"Cannot register task '{name or func.__name__}': Dramatiq not available")
+            return func
+
+        try:
+            import dramatiq
+
+            task_name = name or func.__name__
+
+            # Decorate with dramatiq.actor
+            actor = dramatiq.actor(
+                func,
+                actor_name=task_name,
+                max_retries=max_retries,
+                min_backoff=min_backoff,
+                max_backoff=max_backoff,
+            )
+
+            self._tasks[task_name] = actor
+            log.debug(f"Registered Dramatiq task: {task_name}")
+
+            return actor
+
+        except Exception as e:
+            log.error(f"Failed to register task '{name or func.__name__}': {e}")
+            return func
+
+    def is_available(self) -> bool:
+        """Check if Dramatiq driver is available.
+
+        Returns:
+            True if driver initialized successfully
+        """
+        return self._available
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get queue statistics from Redis.
+
+        Returns:
+            Dictionary with queue statistics
+        """
+        if not self._available or not self._broker:
+            return {"error": "Dramatiq driver not available"}
+
+        try:
+            import redis
+
+            client = redis.from_url(self._redis_url)
+
+            # Get basic stats from Redis
+            stats = {
+                "driver": "dramatiq",
+                "backend": "redis",
+                "available": True,
+                "registered_tasks": list(self._tasks.keys()),
+            }
+
+            # Try to get queue lengths
+            try:
+                keys = client.keys("dramatiq:*:msgs")
+                stats["queues"] = {}
+                for key in keys:
+                    queue_name = key.decode("utf-8").split(":")[1]
+                    length = client.llen(key)
+                    stats["queues"][queue_name] = length
+            except Exception as e:
+                log.debug(f"Could not fetch queue lengths: {e}")
+
+            return stats
+
+        except Exception as e:
+            log.error(f"Failed to get Dramatiq stats: {e}")
+            return {"error": str(e)}
+
+    def get_task_result(self, task_id: Any) -> dict[str, Any]:
+        """Get the result of a task by its ID.
+
+        Note: Dramatiq doesn't have built-in result storage by default.
+        This returns limited information based on message ID.
+
+        Args:
+            task_id: Dramatiq message object
+
+        Returns:
+            Dictionary with task status (limited in Dramatiq without results backend)
+        """
+        if not self._available:
+            return {"status": "error", "error": "Dramatiq driver not available"}
+
+        try:
+            # Dramatiq messages don't have built-in result tracking
+            # unless using Results middleware with a backend
+            return {
+                "status": "pending",
+                "message": "Dramatiq task enqueued. Result tracking requires Results middleware.",
+                "task_id": str(task_id) if task_id else None,
+            }
+        except Exception as e:
+            log.error(f"Failed to get task result: {e}")
+            return {"status": "error", "error": str(e)}
+
+    def get_bulk_results(self, task_ids: list[Any]) -> dict[str, Any]:
+        """Get results for multiple tasks (for bulk feedback: x of y completed).
+
+        Note: Without Results middleware, Dramatiq cannot track task completion.
+        This returns estimated status based on queue inspection.
+
+        Args:
+            task_ids: List of Dramatiq message objects
+
+        Returns:
+            Dictionary with aggregated status (limited without Results backend)
+        """
+        if not self._available:
+            return {"error": "Dramatiq driver not available"}
+
+        total = len(task_ids)
+
+        # Without Results middleware, we can only provide limited feedback
+        return {
+            "total": total,
+            "completed": 0,
+            "failed": 0,
+            "pending": total,  # Assume all pending without result tracking
+            "processing": 0,
+            "tasks": {},
+            "message": "Bulk result tracking requires Results middleware with a backend.",
+            "progress_percentage": 0,
+        }

coati_payroll/queue/drivers/huey_driver.py

@@ -0,0 +1,390 @@
+# Copyright 2025 BMO Soluciones, S.A.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Huey driver for filesystem-based queue processing (fallback mode)."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Callable
+
+from coati_payroll.log import log
+from coati_payroll.queue.driver import QueueDriver
+
+
+class HueyDriver(QueueDriver):
+    """Queue driver using Huey with filesystem backend.
+
+    This driver provides a fallback queue implementation that works
+    without Redis or any database, using only the filesystem for
+    persistence. Suitable for small deployments or development.
+
+    Features:
+    - No external dependencies (Redis, databases)
+    - Persistent queue using filesystem
+    - Thread-safe execution
+    - Multiple workers support (local only)
+
+    Limitations:
+    - Cannot scale horizontally (single-server only)
+    - Lower performance than Dramatiq+Redis
+    """
+
+    def __init__(self, storage_path: str | None = None):
+        """Initialize Huey driver with filesystem backend.
+
+        Args:
+            storage_path: Path to store queue files (default: /var/lib/coati/queue)
+        """
+        self._storage_path = storage_path or self._get_default_storage_path()
+        self._huey = None
+        self._tasks = {}
+        self._available = self._initialize_huey()
+
+    def _get_default_storage_path(self) -> str:
+        """Get default storage path for queue files.
+
+        Ensures proper permissions are available for reading and writing queue files.
+
+        Returns:
+            Path to queue storage directory
+        """
+        # Try to use /var/lib/coati/queue if writable, otherwise use user directory
+        # Try standard system and user directories first (secure locations)
+        paths_to_try = [
+            "/var/lib/coati/queue",
+            os.path.expanduser("~/.local/share/coati-payroll/queue"),
+        ]
+
+        for path in paths_to_try:
+            try:
+                Path(path).mkdir(parents=True, exist_ok=True)
+                # Test read/write permissions
+                test_file = Path(path) / ".test_permissions"
+                test_file.write_text("test")
+                content = test_file.read_text()
+                test_file.unlink()
+
+                if content == "test":
+                    log.info(f"Queue storage path verified with read/write access: {path}")
+                    return path
+            except (OSError, PermissionError) as e:
+                log.debug(f"Cannot use path {path}: {e}")
+                continue
+
+        # Try current working directory as last resort (with warning about security)
+        try:
+            cwd_path = os.path.join(os.getcwd(), ".coati_queue")
+            Path(cwd_path).mkdir(parents=True, exist_ok=True)
+            test_file = Path(cwd_path) / ".test_permissions"
+            test_file.write_text("test")
+            test_file.read_text()
+            test_file.unlink()
+            log.warning(
+                f"Using current working directory for queue storage: {cwd_path}. "
+                f"This may be insecure if running in a public directory. "
+                f"Consider setting COATI_QUEUE_PATH to a secure location."
+            )
+            return cwd_path
+        except (OSError, PermissionError):
+            pass
+
+        # Fallback to temp directory
+        import tempfile
+
+        path = os.path.join(tempfile.gettempdir(), "coati_queue")
+        Path(path).mkdir(parents=True, exist_ok=True)
+        log.warning(
+            f"Using temporary directory for queue storage: {path}. " f"Queue data will be lost on system reboot."
+        )
+        return path
+
+    def _initialize_huey(self) -> bool:
+        """Initialize Huey with filesystem backend.
+
+        Validates read/write permissions before initialization to ensure
+        the process has necessary access to queue files.
+
+        Returns:
+            True if Huey initialized successfully, False otherwise
+        """
+        try:
+            from huey import FileHuey
+
+            # Ensure storage directory exists with proper permissions
+            Path(self._storage_path).mkdir(parents=True, exist_ok=True)
+
+            # Validate read/write permissions
+            test_file = Path(self._storage_path) / ".huey_permissions_test"
+            try:
+                test_file.write_text("permission_test")
+                if test_file.read_text() != "permission_test":
+                    raise PermissionError("Cannot verify read access")
+                test_file.unlink()
+            except (OSError, PermissionError) as e:
+                log.error(
+                    f"Insufficient permissions for queue storage at {self._storage_path}: {e}. "
+                    f"Please ensure the process has read/write access to this directory."
+                )
+                return False
+
+            # Initialize FileHuey with filesystem storage
+            # Note: FileHuey accepts storage_kwargs which are passed to FileStorage
+            self._huey = FileHuey(
+                name="coati_payroll",
+                path=self._storage_path,  # Path for FileStorage
+                immediate=False,  # Don't execute tasks immediately
+                results=True,  # Store results for feedback
+            )
+
+            log.info(
+                f"Huey driver initialized with filesystem storage at {self._storage_path}. "
+                f"Read/write permissions verified."
+            )
+            return True
+
+        except ImportError as e:
+            log.warning(f"Huey not available: {e}")
+            return False
+        except Exception as e:
+            log.error(f"Failed to initialize Huey: {e}")
+            return False
+
+    def enqueue(self, task_name: str, *args: Any, delay: int | None = None, **kwargs: Any) -> Any:
+        """Enqueue a task for background processing.
+
+        Args:
+            task_name: Name of the registered task
+            *args: Positional arguments for the task
+            delay: Optional delay in seconds before execution
+            **kwargs: Keyword arguments for the task
+
+        Returns:
+            Huey result object
+
+        Raises:
+            ValueError: If task is not registered
+            RuntimeError: If driver is not available
+        """
+        if not self._available:
+            raise RuntimeError("Huey driver is not available")
+
+        if task_name not in self._tasks:
+            raise ValueError(f"Task '{task_name}' not registered")
+
+        task = self._tasks[task_name]
+
+        if delay:
+            return task.schedule(args=args, kwargs=kwargs, delay=delay)
+        else:
+            return task(*args, **kwargs)
+
+    def register_task(
+        self,
+        func: Callable,
+        name: str | None = None,
+        max_retries: int = 3,
+        min_backoff: int = 15000,
+        max_backoff: int = 86400000,
+    ) -> Callable:
+        """Register a function as a Huey task.
+
+        Args:
+            func: Function to register
+            name: Optional task name (defaults to function name)
+            max_retries: Maximum retry attempts
+            min_backoff: Minimum backoff in milliseconds (converted to seconds)
+            max_backoff: Maximum backoff in milliseconds (not used by Huey)
+
+        Returns:
+            Huey task that can be called or enqueued
+        """
+        if not self._available or not self._huey:
+            log.warning(f"Cannot register task '{name or func.__name__}': Huey not available")
+            return func
+
+        try:
+            task_name = name or func.__name__
+
+            # Convert milliseconds to seconds for retry delay
+            retry_delay = min_backoff / 1000
+
+            # Decorate with huey.task
+            task = self._huey.task(
+                name=task_name,
+                retries=max_retries,
+                retry_delay=int(retry_delay),
+            )(func)
+
+            self._tasks[task_name] = task
+            log.debug(f"Registered Huey task: {task_name}")
+
+            return task
+
+        except Exception as e:
+            log.error(f"Failed to register task '{name or func.__name__}': {e}")
+            return func
+
+    def is_available(self) -> bool:
+        """Check if Huey driver is available.
+
+        Returns:
+            True if driver initialized successfully
+        """
+        return self._available
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get queue statistics.
+
+        Returns:
+            Dictionary with queue statistics
+        """
+        if not self._available or not self._huey:
+            return {"error": "Huey driver not available"}
+
+        try:
+            stats = {
+                "driver": "huey",
+                "backend": "filesystem",
+                "storage_path": self._storage_path,
+                "available": True,
+                "registered_tasks": list(self._tasks.keys()),
+            }
+
+            # Try to get pending task count
+            try:
+                pending = len(self._huey.pending())
+                stats["pending_tasks"] = pending
+            except Exception as e:
+                log.debug(f"Could not fetch pending tasks: {e}")
+
+            # Try to get scheduled task count
+            try:
+                scheduled = len(self._huey.scheduled())
+                stats["scheduled_tasks"] = scheduled
+            except Exception as e:
+                log.debug(f"Could not fetch scheduled tasks: {e}")
+
+            return stats
+
+        except Exception as e:
+            log.error(f"Failed to get Huey stats: {e}")
+            return {"error": str(e)}
+
+    def get_task_result(self, task_id: Any) -> dict[str, Any]:
+        """Get the result of a task by its ID.
+
+        Huey supports result storage when results=True is enabled.
+
+        Args:
+            task_id: Huey result object (must be callable with no args)
+
+        Returns:
+            Dictionary with task status and result
+        """
+        if not self._available or not self._huey:
+            return {"status": "error", "error": "Huey driver not available"}
+
+        try:
+            # Verify task_id is a Huey result object (callable with no args)
+            if not callable(task_id):
+                return {
+                    "status": "error",
+                    "error": "Invalid task_id: expected Huey result object (callable)",
+                }
+
+            # Try to get the result (non-blocking check)
+            # Huey results raise TaskException or DataStoreTimeout if not ready
+            try:
+                from huey.exceptions import TaskException
+
+                result = task_id()
+                return {
+                    "status": "completed",
+                    "result": result,
+                }
+            except TaskException:
+                # Task is not ready yet or failed
+                return {
+                    "status": "pending",
+                    "message": "Task is still processing",
+                }
+            except Exception as e:
+                # Actual error during result retrieval
+                log.error(f"Error retrieving task result: {e}")
+                return {
+                    "status": "failed",
+                    "error": str(e),
+                }
+
+        except Exception as e:
+            log.error(f"Failed to get task result: {e}")
+            return {"status": "error", "error": str(e)}
+
+    def get_bulk_results(self, task_ids: list[Any]) -> dict[str, Any]:
+        """Get results for multiple tasks (for bulk feedback: x of y completed).
+
+        This provides feedback on bulk operations like parallel payroll processing.
+
+        Args:
+            task_ids: List of Huey result objects
+
+        Returns:
+            Dictionary with aggregated status
+        """
+        if not self._available or not self._huey:
+            return {"error": "Huey driver not available"}
+
+        total = len(task_ids)
+        completed = 0
+        failed = 0
+        pending = 0
+        tasks = {}
+
+        for i, task_id in enumerate(task_ids):
+            try:
+                result_info = self.get_task_result(task_id)
+                status = result_info.get("status", "unknown")
+
+                tasks[f"task_{i}"] = result_info
+
+                if status == "completed":
+                    completed += 1
+                elif status == "error" or status == "failed":
+                    failed += 1
+                else:
+                    pending += 1
+
+            except Exception as e:
+                log.debug(f"Error checking task {i}: {e}")
+                failed += 1
+                tasks[f"task_{i}"] = {"status": "error", "error": str(e)}
+
+        return {
+            "total": total,
+            "completed": completed,
+            "failed": failed,
+            "pending": pending,
+            "processing": 0,  # Huey doesn't distinguish pending from processing
+            "tasks": tasks,
+            "progress_percentage": round((completed / total * 100) if total > 0 else 0, 2),
+        }
+
+    def get_huey_instance(self):
+        """Get the underlying Huey instance for advanced usage.
+
+        Returns:
+            Huey instance or None if not initialized
+        """
+        return self._huey