django-nativemojo 0.1.10__py3-none-any.whl

mojo/apps/tasks/runner.py

@@ -0,0 +1,174 @@
+from importlib import import_module
+from concurrent.futures import ThreadPoolExecutor
+from .manager import TaskManager
+from mojo.tasks import manager
+import os
+from mojo.helpers import logit
+from mojo.helpers import daemon
+from mojo.helpers import paths
+import time
+
+
+class TaskEngine(daemon.Daemon):
+    """
+    The TaskEngine is responsible for managing and executing tasks across different channels.
+    It leverages a thread pool to execute tasks concurrently and uses a task manager to maintain task states.
+    """
+    def __init__(self, channels=["broadcast"], max_workers=5):
+        """
+        Initialize the TaskEngine.
+
+        Args:
+            channels (list): A list of channel names where tasks are queued.
+            max_workers (int, optional): The maximum number of threads available for task execution. Defaults to 5.
+        """
+        super().__init__("taskit", os.path.join(paths.VAR_ROOT, "taskit"))
+        self.manager = manager.TaskManager(channels)
+        self.channels = channels
+        if "broadcast" not in self.channels:
+            self.channels.append("broadcast")
+        self.max_workers = max_workers
+        self.executor = None
+        self.logger = logit.get_logger("taskit", "taskit.log")
+
+    def reset_running_tasks(self):
+        """
+        Reset tasks that are stuck in a running state by moving them back to the pending state.
+        """
+        for channel in self.channels:
+            for task_id in self.manager.get_running_ids(channel):
+                self.logger.info(f"moving task {task_id} from running to pending")
+                self.manager.remove_from_running(channel, task_id)
+                self.manager.add_to_pending(channel, task_id)
+
+    def queue_pending_tasks(self):
+        """
+        Queue all the pending tasks for execution.
+        """
+        for channel in self.channels:
+            for task_id in self.manager.get_pending_ids(channel):
+                self.queue_task(task_id)
+
+    def handle_message(self, message):
+        """
+        Handle incoming messages from the channels, decoding task identifiers and queuing them for execution.
+
+        Args:
+            message (dict): A dictionary with message data containing task information.
+        """
+        self.queue_task(message['data'].decode())
+
+    def on_run_task(self, task_id):
+        """
+        Execute a task based on its identifier by locating the relevant function and executing it.
+
+        Args:
+            task_id (str): The identifier of the task to be executed.
+        """
+        # this is a keep it thread safe with the redis connection
+        tman = TaskManager([])
+        task_data = tman.get_task(task_id)
+        if not task_data:
+            # this task has expired or no longer exists
+            self.logger.info(f"Task {task_id} has expired or no longer exists")
+            tman.remove_from_pending(task_id)
+            return
+        self.logger.info(f"Executing task {task_id}")
+        function_path = task_data.get('function')
+        module_name, func_name = function_path.rsplit('.', 1)
+        module = import_module(module_name)
+        func = getattr(module, func_name)
+        self.manager.remove_from_pending(task_id, task_data.channel)
+        self.manager.add_to_running(task_id, task_data.channel)
+
+        try:
+            task_data.started_at = time.time()
+            func(task_data)
+            task_data.completed_at = time.time()
+            task_data.elapsed_time = task_data.completed_at - task_data.started_at
+            tman.save_task(task_data)
+            tman.add_to_completed(task_data)
+            self.logger.info(f"Task {task_id} completed after {task_data.elapsed_time} seconds")
+        except Exception as e:
+            self.logger.error(f"Error executing task {task_id}: {str(e)}")
+            tman.add_to_errors(task_data, str(e))
+        finally:
+            tman.remove_from_running(task_id, task_data.channel)
+
+    def queue_task(self, task_id):
+        """
+        Submit a task for execution in the thread pool.
+
+        Args:
+            task_id (str): The identifier of the task to be queued.
+        """
+        self.logger.info(f"adding task {task_id}")
+        self.executor.submit(self.on_run_task, task_id)
+
+    def wait_for_all_tasks_to_complete(self, timeout=30):
+        """
+        Wait for all tasks submitted to the executor to complete.
+        """
+        self.executor.shutdown(wait=True, timeout=timeout)
+        # Check if there are still active threads
+        active_threads = [thread for thread in self.executor._threads if thread.is_alive()]
+        if active_threads:
+            self.logger.warning(f"shutdown issue, {len(active_threads)} tasks exceeded timeout")
+            self.executor.shutdown(wait=False)  # Stop accepting new tasks
+
+    def start_listening(self):
+        """
+        Listen for messages on the subscribed channels and handle them as they arrive.
+        """
+        self.logger.info("starting with channels...", self.channels)
+        self.reset_running_tasks()
+        self.queue_pending_tasks()
+        pubsub = self.manager.redis.pubsub()
+        channel_keys = {self.manager.get_channel_key(channel): self.handle_message for channel in self.channels}
+        pubsub.subscribe(**channel_keys)
+        for message in pubsub.listen():
+            if not self.running:
+                self.logger.info("shutting down, waiting for tasks to complete")
+                self.wait_for_all_tasks_to_complete()
+                self.logger.info("shutdown complete")
+                return
+            if message['type'] != 'message':
+                continue
+            self.handle_message(message)
+
+    def run(self):
+        self.executor = ThreadPoolExecutor(max_workers=self.max_workers)
+        self.start_listening()
+
+
+# HELPERS FOR RUNNING VIA CLI
+def get_args():
+    """
+    Setup the argument parser for command-line interface.
+
+    Returns:
+        Namespace: Parsed command-line arguments.
+    """
+    import argparse
+    parser = argparse.ArgumentParser(description="TaskEngine Background Service")
+    parser.add_argument("--start", action="store_true", help="Start the daemon")
+    parser.add_argument("--stop", action="store_true", help="Stop the daemon")
+    parser.add_argument("--foreground", "-f", action="store_true", help="Run in foreground mode")
+    parser.add_argument("-v", "--verbose", action="store_true",
+                        help="Enable verbose logging")
+    return parser, parser.parse_args()
+
+
+def main():
+    from mojo.helpers.settings import settings
+    parser, args = get_args()
+    daemon = TaskEngine(settings.TASKIT_CHANNELS)
+    if args.start:
+        daemon.start()
+    elif args.stop:
+        daemon.stop()
+    elif args.foreground:
+        print("Running in foreground mode...")
+        daemon.run()
+    else:
+        parser.print_help()

mojo/apps/tasks/tq_handlers.py

@@ -0,0 +1,14 @@
+from mojo.helpers import logit
+import time
+
+logger = logit.get_logger("ti_example", "ti_example.log")
+
+def run_example_task(task):
+    logger.info("Running example task with data", task)
+    time.sleep(task.data.get("duration", 5))
+
+
+def run_error_task(task):
+    logger.info("Running error task with data", task)
+    time.sleep(2)
+    raise Exception("Example error")

mojo/decorators/__init__.py

@@ -0,0 +1,3 @@
+from .http import *
+from .validate import *
+from .auth import *

mojo/decorators/auth.py

@@ -0,0 +1,25 @@
+from functools import wraps
+import mojo.errors
+
+def requires_perms(*required_perms):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(request, *args, **kwargs):
+            if not request.user.is_authenticated:
+                raise mojo.errors.PermissionDeniedException()
+            if not request.user.has_permission(required_perms):
+                raise mojo.errors.PermissionDeniedException()
+            return func(request, *args, **kwargs)
+        return wrapper
+    return decorator
+
+
+def requires_auth():
+    def decorator(func):
+        @wraps(func)
+        def wrapper(request, *args, **kwargs):
+            if not request.user.is_authenticated:
+                raise mojo.errors.PermissionDeniedException()
+            return func(request, *args, **kwargs)
+        return wrapper
+    return decorator

mojo/decorators/cron.py

@@ -0,0 +1,31 @@
+from typing import Callable
+
+def schedule(minutes: str = '*', hours: str = '*', days: str = '*',
+                   months: str = '*', weekdays: str = '*') -> Callable:
+    """
+    A decorator to schedule functions based on cron syntax.
+
+    Args:
+        minutes (str): The minutes argument for the cron schedule (default is '*').
+        hours (str): The hours argument for the cron schedule (default is '*').
+        days (str): The days of the month argument for the cron schedule (default is '*').
+        months (str): The months argument for the cron schedule (default is '*').
+        weekdays (str): The days of the week argument for the cron schedule (default is '*').
+
+    Returns:
+        Callable: The decorated function.
+    """
+    def decorator(func: Callable) -> Callable:
+        if not hasattr(decorator, 'scheduled_functions'):
+            decorator.scheduled_functions = []
+        cron_spec = {
+            'func': func,
+            'minutes': minutes,
+            'hours': hours,
+            'days': days,
+            'months': months,
+            'weekdays': weekdays
+        }
+        decorator.scheduled_functions.append(cron_spec)
+        return func
+    return decorator

mojo/decorators/http.py

@@ -0,0 +1,132 @@
+import sys
+import traceback
+from mojo.helpers.settings import settings
+from mojo.helpers import modules as jm
+from mojo.helpers import logit
+import mojo.errors
+from django.urls import path, re_path
+# from django.http import JsonResponse
+from mojo.helpers.response import JsonResponse
+from functools import wraps
+from mojo.helpers import modules
+from mojo.models import rest
+from mojo.apps import metrics
+
+logger = logit.get_logger("error", "error.log")
+# logger.info("created")
+
+# Global registry for REST routes
+REGISTERED_URLS = {}
+URLPATTERN_METHODS = {}
+MOJO_API_MODULE = settings.get("MOJO_API_MODULE", "api")
+MOJO_APPEND_SLASH = settings.get("MOJO_APPEND_SLASH", False)
+
+API_METRICS = settings.get("API_METRICS", False)
+API_METRICS_GRANULARITY = settings.get("API_METRICS_GRANULARITY", "days")
+
+
+def dispatcher(request, *args, **kwargs):
+    """
+    Dispatches incoming requests to the appropriate registered URL method.
+    """
+    rest.ACTIVE_REQUEST = request
+    key = kwargs.pop('__mojo_rest_root_key__', None)
+    if "group" in request.DATA:
+        request.group = modules.get_model_instance("account", "Group", int(request.DATA.group))
+    method_key = f"{key}__{request.method}"
+    if method_key not in URLPATTERN_METHODS:
+        method_key = f"{key}__ALL"
+    if method_key in URLPATTERN_METHODS:
+        return dispatch_error_handler(URLPATTERN_METHODS[method_key])(request, *args, **kwargs)
+    return JsonResponse({"error": "Endpoint not found", "code": 404}, status=404)
+
+
+def dispatch_error_handler(func):
+    """
+    Decorator to catch and handle errors.
+    It logs exceptions and returns appropriate HTTP responses.
+    """
+    @wraps(func)
+    def wrapper(request, *args, **kwargs):
+        try:
+            if API_METRICS:
+                metrics.record("api_calls", category="mojo_api", min_granularity=API_METRICS_GRANULARITY)
+            return func(request, *args, **kwargs)
+        except mojo.errors.MojoException as err:
+            if API_METRICS:
+                metrics.record("api_errors", category="mojo_api", min_granularity=API_METRICS_GRANULARITY)
+            return JsonResponse({"error": err.reason, "code": err.code}, status=err.status)
+        except ValueError as err:
+            if API_METRICS:
+                metrics.record("api_errors", category="mojo_api", min_granularity=API_METRICS_GRANULARITY)
+            logger.exception(f"ValueErrror: {str(err)}, Path: {request.path}, IP: {request.META.get('REMOTE_ADDR')}")
+            return JsonResponse({"error": str(err), "code": 555 }, status=500)
+        except Exception as err:
+            if API_METRICS:
+                metrics.record("api_errors", category="mojo_api", min_granularity=API_METRICS_GRANULARITY)
+            # logger.exception(f"Unhandled REST Exception: {request.path}")
+            logger.exception(f"Error: {str(err)}, Path: {request.path}, IP: {request.META.get('REMOTE_ADDR')}")
+            return JsonResponse({"error": str(err) }, status=500)
+
+    return wrapper
+
+
+def _register_route(method="ALL"):
+    """
+    Decorator to automatically register a Django view for a specific HTTP method.
+    Supports defining a custom pattern inside the decorator.
+
+    :param method: The HTTP method (GET, POST, etc.).
+    """
+    def decorator(pattern=None, docs=None):
+        def wrapper(view_func):
+            module = jm.get_root_module(view_func)
+            if not module:
+                print("!!!!!!!")
+                print(sys._getframe(2).f_code.co_filename)
+                raise RuntimeError(f"Could not determine module for {view_func.__name__}")
+
+            # Ensure `urlpatterns` exists in the calling module
+            if not hasattr(module, 'urlpatterns'):
+                module.urlpatterns = []
+
+            # If no pattern is provided, use the function name as the pattern
+            if pattern is None:
+                pattern_used = f"{view_func.__name__}"
+            else:
+                pattern_used = pattern
+
+            if MOJO_APPEND_SLASH:
+                pattern_used = pattern if pattern_used.endswith("/") else f"{pattern_used}/"
+            # Register view in URL mapping
+            app_name = module.__name__.split(".")[-1]
+            # print(f"{module.__name__}.urlpatterns")
+            root_key = f"{app_name}__{pattern_used}"
+            key = f"{root_key}__{method}"
+            # print(f"{app_name} -> {pattern_used} -> {key}")
+            URLPATTERN_METHODS[key] = view_func
+
+            # Determine whether to use path() or re_path()
+            url_func = path if not (pattern_used.startswith("^") or pattern_used.endswith("$")) else re_path
+
+            # Add to `urlpatterns`
+            module.urlpatterns.append(url_func(
+                pattern_used, dispatcher,
+                kwargs={
+                    "__mojo_rest_root_key__": root_key
+                }))
+            # Attach metadata
+            view_func.__app_module_name__ = module.__name__
+            view_func.__app_name__ = app_name
+            view_func.__url__ = (method, pattern_used)
+            view_func.__docs__ = docs or {}
+            return view_func
+        return wrapper
+    return decorator
+
+# Public-facing URL decorators
+URL = _register_route()
+GET = _register_route("GET")
+POST = _register_route("POST")
+PUT = _register_route("PUT")
+DELETE = _register_route("DELETE")

mojo/decorators/validate.py

@@ -0,0 +1,14 @@
+from functools import wraps
+import mojo.errors
+
+def requires_params(*required_params):
+    def decorator(func):
+        @wraps(func)
+        def wrapper(request, *args, **kwargs):
+            missing_params = [param for param in required_params if param not in request.DATA]
+            if missing_params:
+                str_params = ', '.join(missing_params)
+                raise mojo.errors.ValueException(f"missing required parameters: {str_params}")
+            return func(request, *args, **kwargs)
+        return wrapper
+    return decorator

mojo/errors.py

@@ -0,0 +1,88 @@
+class MojoException(Exception):
+    """
+    Base exception class for Mojo-related errors.
+
+    Attributes:
+        reason (str): The reason for the exception.
+        code (int): The error code associated with the exception.
+        status (int, optional): The HTTP status code. Defaults to None.
+    """
+
+    def __init__(self, reason, code, status=500):
+        """
+        Initialize a MojoException instance.
+
+        Args:
+            reason (str): The reason for the exception.
+            code (int): The error code associated with the exception.
+            status (int, optional): The HTTP status code. Defaults to None.
+        """
+        super().__init__(reason)
+        self.reason = reason
+        self.code = code
+        self.status = status
+
+
+class ValueException(MojoException):
+    """
+    Exception raised for REST API value errors.
+
+    Attributes:
+        reason (str): The reason for the exception. Defaults to 'REST API Error'.
+        code (int): The error code associated with the exception. Defaults to 500.
+        status (int, optional): The HTTP status code. Defaults to 500.
+    """
+
+    def __init__(self, reason='REST API Error', code=400, status=400):
+        """
+        Initialize a RestErrorException instance.
+
+        Args:
+            reason (str, optional): The reason for the exception. Defaults to 'REST API Error'.
+            code (int, optional): The error code associated with the exception. Defaults to 500.
+            status (int, optional): The HTTP status code. Defaults to 500.
+        """
+        super().__init__(reason, code, status)
+
+
+class PermissionDeniedException(MojoException):
+    """
+    Exception raised for permission denied errors.
+
+    Attributes:
+        reason (str): The reason for the exception. Defaults to 'Permission Denied'.
+        code (int): The error code associated with the exception. Defaults to 403.
+        status (int, optional): The HTTP status code. Defaults to 403.
+    """
+
+    def __init__(self, reason='Permission Denied', code=403, status=403):
+        """
+        Initialize a PermissionDeniedException instance.
+
+        Args:
+            reason (str, optional): The reason for the exception. Defaults to 'Permission Denied'.
+            code (int, optional): The error code associated with the exception. Defaults to 403.
+            status (int, optional): The HTTP status code. Defaults to 403.
+        """
+        super().__init__(reason, code, status)
+
+class RestErrorException(MojoException):
+    """
+    Exception raised for REST API errors.
+
+    Attributes:
+        reason (str): The reason for the exception. Defaults to 'REST API Error'.
+        code (int): The error code associated with the exception. Defaults to 500.
+        status (int, optional): The HTTP status code. Defaults to 500.
+    """
+
+    def __init__(self, reason='REST API Error', code=500, status=500):
+        """
+        Initialize a RestErrorException instance.
+
+        Args:
+            reason (str, optional): The reason for the exception. Defaults to 'REST API Error'.
+            code (int, optional): The error code associated with the exception. Defaults to 500.
+            status (int, optional): The HTTP status code. Defaults to 500.
+        """
+        super().__init__(reason, code, status)

mojo/helpers/aws/client.py

@@ -0,0 +1,8 @@
+import boto3
+
+def get_session(access_key, secret_key, region):
+    return boto3.Session(
+        aws_access_key_id=access_key,
+        aws_secret_access_key=secret_key,
+        region_name=region,
+    )