apisec-code-bolt 0.1.0__py3-none-any.whl

apisec_code_bolt/frameworks/python/celery_plugin.py

@@ -0,0 +1,393 @@
+"""
+Task framework plugin for Celery, Dramatiq, Huey, and RQ.
+
+Detects async task entry points from decorator patterns across the four
+most widely-used Python task/worker frameworks.
+
+Supported constructs — Celery (5.x):
+- @app.task / @app.task(name="...", bind=True, ...)
+- @shared_task / @shared_task(name="...", ...)
+- @celery.task (alternate Celery instance naming)
+- Class-based tasks (subclasses of celery.Task) — NOT yet supported
+
+Supported constructs — Dramatiq (2.x):
+- @dramatiq.actor / @dramatiq.actor(queue_name="...", actor_name="...", ...)
+
+Supported constructs — Huey (2.x):
+- @huey.task() / @huey.periodic_task(...)
+- Alternate instance names: @queue.task(), @queue.periodic_task()
+
+Supported constructs — RQ (1.x):
+- @job(queue, connection, ...) from rq.decorators
+
+Not supported (yet):
+- Celery class-based tasks (subclass of celery.Task with run() method)
+- Celery canvas primitives (chain, group, chord as entry points)
+- Dynamic task registration (app.register_task)
+- Dramatiq middleware hooks
+- Huey pre/post-execute hooks
+- RQ queue.enqueue() calls as entry points
+"""
+
+from __future__ import annotations
+
+from typing import Any, ClassVar
+
+from ...core.types import (
+    Confidence,
+    Framework,
+    HttpMethod,
+    Language,
+    ParameterLocation,
+)
+from ...parsing.base import ParsedDecorator, ParsedFile, ParsedFunction
+from ..base import (
+    BaseFrameworkPlugin,
+    ExtractedAuthDependency,
+    ExtractedAuthScheme,
+    ExtractedDependency,
+    ExtractedMiddleware,
+    ExtractedParameter,
+    ExtractedResponse,
+    ExtractedRoute,
+    FrameworkPluginRegistry,
+)
+
+# ── Import detection ────────────────────────────────────────────────────────
+
+CELERY_IMPORTS = frozenset(
+    {
+        "celery",
+        "celery.app",
+        "celery.shared_task",
+    }
+)
+
+DRAMATIQ_IMPORTS = frozenset(
+    {
+        "dramatiq",
+    }
+)
+
+HUEY_IMPORTS = frozenset(
+    {
+        "huey",
+    }
+)
+
+RQ_IMPORTS = frozenset(
+    {
+        "rq",
+        "rq.decorators",
+    }
+)
+
+ALL_TASK_IMPORTS = CELERY_IMPORTS | DRAMATIQ_IMPORTS | HUEY_IMPORTS | RQ_IMPORTS
+
+
+# ── Decorator name sets ─────────────────────────────────────────────────────
+
+# Celery: @app.task, @celery_app.task, @shared_task
+CELERY_TASK_DECORATORS: set[str] = {
+    "task",  # @app.task / @celery.task (attribute form)
+    "shared_task",  # @shared_task (bare import from celery)
+    "celery.shared_task",  # @celery.shared_task (qualified)
+}
+
+# Dramatiq: @dramatiq.actor
+DRAMATIQ_ACTOR_DECORATORS: set[str] = {
+    "actor",  # @broker.actor (attribute) or @dramatiq.actor (qualified)
+    "dramatiq.actor",  # @dramatiq.actor
+}
+
+# Huey: @huey.task(), @huey.periodic_task()
+HUEY_TASK_DECORATORS: set[str] = {
+    "task",  # @huey.task() (attribute)
+    "periodic_task",  # @huey.periodic_task()
+}
+
+# RQ: @job(queue, connection)
+RQ_JOB_DECORATORS: set[str] = {
+    "job",  # @job (from rq.decorators import job)
+    "rq.decorators.job",  # qualified import
+}
+
+ALL_TASK_DECORATORS = (
+    CELERY_TASK_DECORATORS | DRAMATIQ_ACTOR_DECORATORS | HUEY_TASK_DECORATORS | RQ_JOB_DECORATORS
+)
+
+
+# ── Which framework set detected? ──────────────────────────────────────────
+
+
+class _DetectedFramework:
+    """Tracks which task framework(s) are present in a file."""
+
+    __slots__ = ("celery", "dramatiq", "huey", "rq")
+
+    def __init__(self) -> None:
+        self.celery = False
+        self.dramatiq = False
+        self.huey = False
+        self.rq = False
+
+    @property
+    def any(self) -> bool:
+        return self.celery or self.dramatiq or self.huey or self.rq
+
+
+def _detect_task_frameworks(parsed_file: ParsedFile) -> _DetectedFramework:
+    found = _DetectedFramework()
+    for imp in parsed_file.imports:
+        module = imp.module or ""
+        mod = module.split(".")[0]
+        if mod == "celery" or module.endswith(".celery"):
+            # Catches direct `from celery import ...` AND the common Django
+            # convention of `from myapp.celery import app` where the app's
+            # Celery instance lives in a local `celery.py` module.
+            found.celery = True
+        elif mod == "dramatiq":
+            found.dramatiq = True
+        elif mod == "huey":
+            found.huey = True
+        elif mod == "rq":
+            found.rq = True
+        for name in imp.names:
+            if name == "shared_task":
+                found.celery = True
+            elif name == "job" and imp.module and "rq" in imp.module:
+                found.rq = True
+    return found
+
+
+# ── Plugin ──────────────────────────────────────────────────────────────────
+
+
+class CeleryPlugin(BaseFrameworkPlugin):
+    """Plugin for Celery, Dramatiq, Huey, and RQ task frameworks."""
+
+    FRAMEWORK: ClassVar[Framework] = Framework.CELERY
+    LANGUAGE: ClassVar[Language] = Language.PYTHON
+    DETECTION_IMPORTS: ClassVar[frozenset[str]] = ALL_TASK_IMPORTS
+
+    # ------------------------------------------------------------------
+    # Detection
+    # ------------------------------------------------------------------
+
+    def detect(self, parsed_file: ParsedFile) -> bool:
+        return _detect_task_frameworks(parsed_file).any
+
+    # ------------------------------------------------------------------
+    # Route (task) extraction
+    # ------------------------------------------------------------------
+
+    def extract_routes(
+        self,
+        parsed_file: ParsedFile,
+        context: Any = None,
+    ) -> list[ExtractedRoute]:
+        fw = _detect_task_frameworks(parsed_file)
+        routes: list[ExtractedRoute] = []
+
+        for func in parsed_file.functions:
+            route = self._try_extract_task(func, fw)
+            if route:
+                routes.append(route)
+
+        for cls in parsed_file.classes:
+            for method in cls.methods:
+                route = self._try_extract_task(method, fw)
+                if route:
+                    routes.append(route)
+
+        return routes
+
+    # ------------------------------------------------------------------
+    # Stubs — task frameworks don't define HTTP auth/deps/middleware
+    # ------------------------------------------------------------------
+
+    def extract_dependencies(self, parsed_file: ParsedFile) -> list[ExtractedDependency]:
+        return []
+
+    def extract_auth_schemes(self, parsed_file: ParsedFile) -> list[ExtractedAuthScheme]:
+        return []
+
+    def extract_auth_dependencies(
+        self,
+        parsed_file: ParsedFile,
+        known_scheme_names: set[str] | None = None,
+        **kwargs: Any,
+    ) -> list[ExtractedAuthDependency]:
+        return []
+
+    def extract_middleware(self, parsed_file: ParsedFile) -> list[ExtractedMiddleware]:
+        return []
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _try_extract_task(
+        self,
+        func: ParsedFunction,
+        fw: _DetectedFramework,
+    ) -> ExtractedRoute | None:
+        """If *func* has a task decorator, return an ExtractedRoute with kind="task"."""
+        dec, origin = self._find_task_decorator(func, fw)
+        if dec is None:
+            return None
+
+        task_name = self._resolve_task_name(dec, func, origin)
+        params = self._extract_parameters(func, origin)
+        kind = "scheduled" if self._is_periodic(dec) else "task"
+
+        return ExtractedRoute(
+            kind=kind,
+            method=HttpMethod.POST,
+            path=task_name,
+            handler_function=func.qualified_name,
+            handler_location=func.location,
+            path_params=[],
+            query_params=params,
+            header_params=[],
+            cookie_params=[],
+            body=None,
+            response=ExtractedResponse(),
+            confidence=Confidence.HIGH,
+        )
+
+    def _find_task_decorator(
+        self,
+        func: ParsedFunction,
+        fw: _DetectedFramework,
+    ) -> tuple[ParsedDecorator | None, str]:
+        """Return (decorator, origin_framework) if a task decorator is found."""
+        for dec in func.decorators:
+            base = self._decorator_base_name(dec)
+
+            if fw.celery and base in CELERY_TASK_DECORATORS:
+                return dec, "celery"
+            if fw.dramatiq and base in DRAMATIQ_ACTOR_DECORATORS:
+                return dec, "dramatiq"
+            if fw.huey and base in HUEY_TASK_DECORATORS:
+                return dec, "huey"
+            if fw.rq and base in RQ_JOB_DECORATORS:
+                return dec, "rq"
+
+        return None, ""
+
+    def _resolve_task_name(
+        self,
+        dec: ParsedDecorator,
+        func: ParsedFunction,
+        origin: str,
+    ) -> str:
+        """Derive the canonical task name.
+
+        Priority: explicit ``name=`` kwarg > first positional string arg >
+        function qualified name.
+        """
+        name_kwarg_key = "actor_name" if origin == "dramatiq" else "name"
+        explicit = self._extract_decorator_arg(dec, name_kwarg_key)
+        if explicit and isinstance(explicit, str):
+            return explicit
+
+        if dec.positional_args:
+            first = dec.positional_args[0]
+            if isinstance(first, str):
+                return first
+
+        return func.qualified_name.full
+
+    @staticmethod
+    def _is_periodic(dec: ParsedDecorator) -> bool:
+        """True if the decorator indicates a periodic/scheduled task (Huey)."""
+        name = dec.name or ""
+        return "periodic" in name.lower()
+
+    def _extract_parameters(
+        self,
+        func: ParsedFunction,
+        origin: str,
+    ) -> list[ExtractedParameter]:
+        """Extract task argument parameters from the function signature.
+
+        Celery ``bind=True`` injects ``self`` as the first parameter — skip it.
+        """
+        params: list[ExtractedParameter] = []
+        skip_first = False
+
+        if origin == "celery":
+            for dec in func.decorators:
+                bind_val = self._extract_decorator_arg(dec, "bind")
+                if bind_val in (True, "True", "true"):
+                    skip_first = True
+                    break
+
+        skipped = False
+        for p in func.parameters:
+            if p.name in ("self", "cls"):
+                continue
+            if skip_first and not skipped:
+                skipped = True
+                continue
+            if p.is_variadic or p.is_keyword_variadic:
+                continue
+
+            params.append(
+                ExtractedParameter(
+                    name=p.name,
+                    location=ParameterLocation.TASK_ARGUMENT,
+                    type_annotation=p.type_annotation,
+                    required=p.default_value is None,
+                    default_value=p.default_value,
+                )
+            )
+
+        return params
+
+    @staticmethod
+    def _decorator_base_name(dec: ParsedDecorator) -> str:
+        """Normalise a decorator to a lookup key.
+
+        Handles:
+        - ``celery.shared_task`` (qualified)
+        - ``shared_task`` (bare import)
+        - ``app.task`` (attribute on Celery instance)
+        - ``dramatiq.actor`` (qualified)
+        - ``huey.task`` / ``huey.periodic_task`` (attribute on Huey instance)
+        - ``job`` (bare RQ import)
+        """
+        if dec.qualified_name:
+            qn = dec.qualified_name.full
+            parts = qn.rsplit(".", 1)
+            if len(parts) == 2:
+                mod, attr = parts
+                if mod in ("celery", "dramatiq", "rq.decorators"):
+                    return f"{mod}.{attr}"
+                if attr in (
+                    "task",
+                    "shared_task",
+                    "actor",
+                    "periodic_task",
+                    "job",
+                ):
+                    return attr
+            return qn
+
+        name = dec.name
+        if "." in name:
+            _obj, _, attr = name.rpartition(".")
+            if attr in (
+                "task",
+                "shared_task",
+                "actor",
+                "periodic_task",
+                "job",
+            ):
+                return attr
+        return name
+
+
+_celery_plugin = CeleryPlugin()
+FrameworkPluginRegistry.register(_celery_plugin)