codegraph-engine 2.2.1__py3-none-any.whl → 3.0.0__py3-none-any.whl

codegraph/async_queue/detectors/celery.py

@@ -0,0 +1,160 @@
+"""Celery asynchronous task and canvas detector.
+
+Detects:
+  - Workers / Consumers: @app.task, @celery.task, @shared_task, @celery_app.task
+  - Publishers / Producers: task.delay(), task.apply_async(), app.send_task()
+"""
+from __future__ import annotations
+
+import ast
+from typing import Any
+
+from codegraph.async_queue.models import (
+    AsyncConsumerRecord,
+    AsyncMessagingSystem,
+    AsyncProducerRecord,
+)
+
+_CELERY_TASK_DECORATORS = {
+    "task",
+    "shared_task",
+}
+
+
+def detect_celery(
+    tree: ast.AST,
+    content: str,
+    file_path: str,
+    module_name: str,
+) -> tuple[list[AsyncProducerRecord], list[AsyncConsumerRecord]]:
+    producers: list[AsyncProducerRecord] = []
+    consumers: list[AsyncConsumerRecord] = []
+    lines = content.splitlines()
+
+    def get_line_snippet(lineno: int) -> str:
+        if 1 <= lineno <= len(lines):
+            return lines[lineno - 1].strip()
+        return ""
+
+    # Map of function names that are Celery tasks in this file
+    local_tasks: dict[str, str] = {}  # fn_name -> task_name
+
+    # 1. Detect consumers (task definitions)
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            for dec in node.decorator_list:
+                dec_func = dec.func if isinstance(dec, ast.Call) else dec
+                dec_id = getattr(dec_func, "id", "")
+                dec_attr = getattr(dec_func, "attr", "")
+
+                if dec_id in _CELERY_TASK_DECORATORS or dec_attr in _CELERY_TASK_DECORATORS:
+                    # Found a Celery task definition
+                    task_name = f"{module_name}.{node.name}"
+                    metadata: dict[str, Any] = {}
+
+                    if isinstance(dec, ast.Call):
+                        for kw in dec.keywords:
+                            if kw.arg == "name" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                                task_name = kw.value.value
+                            elif kw.arg == "queue" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                                metadata["queue"] = kw.value.value
+                            elif kw.arg == "max_retries" and isinstance(kw.value, ast.Constant):
+                                metadata["max_retries"] = kw.value.value
+                            elif kw.arg == "autoretry_for":
+                                metadata["autoretry"] = True
+                            elif kw.arg == "retry_backoff":
+                                metadata["retry_backoff"] = True
+
+                    local_tasks[node.name] = task_name
+                    canonical_id = f"{module_name}.{node.name}"
+                    queue_id = f"celery:{task_name}"
+
+                    consumers.append(
+                        AsyncConsumerRecord(
+                            system=AsyncMessagingSystem.CELERY,
+                            destination=task_name,
+                            queue_id=queue_id,
+                            handler_canonical_id=canonical_id,
+                            file_path=file_path,
+                            line=node.lineno,
+                            handler_snippet=get_line_snippet(node.lineno),
+                            confidence="HIGH",
+                            metadata=metadata,
+                        )
+                    )
+
+    # 2. Detect producers (task invocations: .delay, .apply_async, app.send_task)
+    class CallVisitor(ast.NodeVisitor):
+        def __init__(self) -> None:
+            self.current_function: str = f"{module_name}.<module>"
+
+        def visit_FunctionDef(self, fn: ast.FunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_AsyncFunctionDef(self, fn: ast.AsyncFunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_Call(self, node: ast.Call) -> None:
+            if isinstance(node.func, ast.Attribute):
+                method_name = node.func.attr
+                # Case A: task.delay(...) or task.apply_async(...)
+                if method_name in ("delay", "apply_async"):
+                    caller_val = node.func.value
+                    task_target_name = ""
+                    if isinstance(caller_val, ast.Name):
+                        task_target_name = caller_val.id
+                    elif isinstance(caller_val, ast.Attribute):
+                        task_target_name = caller_val.attr
+
+                    if task_target_name:
+                        resolved_dest = local_tasks.get(task_target_name, task_target_name)
+                        meta: dict[str, Any] = {"method": method_name}
+                        if method_name == "apply_async":
+                            for kw in node.keywords:
+                                if kw.arg in ("countdown", "eta", "queue", "retry"):
+                                    meta[kw.arg] = True
+
+                        producers.append(
+                            AsyncProducerRecord(
+                                system=AsyncMessagingSystem.CELERY,
+                                destination=resolved_dest,
+                                queue_id=f"celery:{resolved_dest}",
+                                caller_canonical_id=self.current_function,
+                                file_path=file_path,
+                                line=node.lineno,
+                                call_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata=meta,
+                            )
+                        )
+
+                # Case B: app.send_task("task_name", ...)
+                elif method_name == "send_task" and node.args:
+                    first_arg = node.args[0]
+                    if isinstance(first_arg, ast.Constant) and isinstance(first_arg.value, str):
+                        dest = first_arg.value
+                        producers.append(
+                            AsyncProducerRecord(
+                                system=AsyncMessagingSystem.CELERY,
+                                destination=dest,
+                                queue_id=f"celery:{dest}",
+                                caller_canonical_id=self.current_function,
+                                file_path=file_path,
+                                line=node.lineno,
+                                call_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata={"method": "send_task"},
+                            )
+                        )
+
+            self.generic_visit(node)
+
+    CallVisitor().visit(tree)
+
+    return producers, consumers

codegraph/async_queue/detectors/kafka.py

@@ -0,0 +1,168 @@
+"""Apache Kafka producer and consumer detector.
+
+Detects:
+  - Publishers / Producers: producer.send(topic, ...), producer.produce(topic, ...)
+  - Consumers / Workers: @consumer(topic), @subscriber(topic), KafkaConsumer(topic), consumer.subscribe([topic])
+"""
+from __future__ import annotations
+
+import ast
+from typing import Any
+
+from codegraph.async_queue.models import (
+    AsyncConsumerRecord,
+    AsyncMessagingSystem,
+    AsyncProducerRecord,
+)
+
+_KAFKA_CONSUMER_DECORATORS = {
+    "consumer",
+    "kafka_consumer",
+    "subscriber",
+    "consume",
+    "agent",
+}
+
+
+def detect_kafka(
+    tree: ast.AST,
+    content: str,
+    file_path: str,
+    module_name: str,
+) -> tuple[list[AsyncProducerRecord], list[AsyncConsumerRecord]]:
+    producers: list[AsyncProducerRecord] = []
+    consumers: list[AsyncConsumerRecord] = []
+    lines = content.splitlines()
+
+    def get_line_snippet(lineno: int) -> str:
+        if 1 <= lineno <= len(lines):
+            return lines[lineno - 1].strip()
+        return ""
+
+    # 1. Detect consumers via decorators or subscriber calls
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            for dec in node.decorator_list:
+                dec_func = dec.func if isinstance(dec, ast.Call) else dec
+                dec_name = getattr(dec_func, "id", "") or getattr(dec_func, "attr", "")
+
+                if dec_name.lower() in _KAFKA_CONSUMER_DECORATORS:
+                    topic = ""
+                    if isinstance(dec, ast.Call):
+                        if dec.args and isinstance(dec.args[0], ast.Constant) and isinstance(dec.args[0].value, str):
+                            topic = dec.args[0].value
+                        for kw in dec.keywords:
+                            if kw.arg in ("topic", "topics") and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                                topic = kw.value.value
+
+                    if topic:
+                        consumers.append(
+                            AsyncConsumerRecord(
+                                system=AsyncMessagingSystem.KAFKA,
+                                destination=topic,
+                                queue_id=f"kafka:{topic}",
+                                handler_canonical_id=f"{module_name}.{node.name}",
+                                file_path=file_path,
+                                line=node.lineno,
+                                handler_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata={"decorator": dec_name},
+                            )
+                        )
+
+    # 2. Detect producers & programmatic consumers inside function scopes
+    class KafkaCallVisitor(ast.NodeVisitor):
+        def __init__(self) -> None:
+            self.current_function: str = f"{module_name}.<module>"
+
+        def visit_FunctionDef(self, fn: ast.FunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_AsyncFunctionDef(self, fn: ast.AsyncFunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_Call(self, node: ast.Call) -> None:
+            # Check for producer.send(topic, ...) or producer.produce(topic, ...)
+            if isinstance(node.func, ast.Attribute) and node.func.attr in ("send", "produce"):
+                topic = ""
+                meta: dict[str, Any] = {"method": node.func.attr}
+                if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
+                    topic = node.args[0].value
+                for kw in node.keywords:
+                    if kw.arg == "topic" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                        topic = kw.value.value
+                    elif kw.arg in ("key", "partition", "headers"):
+                        meta[kw.arg] = True
+
+                if topic:
+                    producers.append(
+                        AsyncProducerRecord(
+                            system=AsyncMessagingSystem.KAFKA,
+                            destination=topic,
+                            queue_id=f"kafka:{topic}",
+                            caller_canonical_id=self.current_function,
+                            file_path=file_path,
+                            line=node.lineno,
+                            call_snippet=get_line_snippet(node.lineno),
+                            confidence="HIGH",
+                            metadata=meta,
+                        )
+                    )
+
+            # Check for consumer.subscribe([topic, ...])
+            elif isinstance(node.func, ast.Attribute) and node.func.attr == "subscribe":
+                topic_list: list[str] = []
+                if node.args and isinstance(node.args[0], (ast.List, ast.Tuple)):
+                    for elt in node.args[0].elts:
+                        if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
+                            topic_list.append(elt.value)
+                for kw in node.keywords:
+                    if kw.arg == "topics" and isinstance(kw.value, (ast.List, ast.Tuple)):
+                        for elt in kw.value.elts:
+                            if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
+                                topic_list.append(elt.value)
+
+                for t in topic_list:
+                    consumers.append(
+                        AsyncConsumerRecord(
+                            system=AsyncMessagingSystem.KAFKA,
+                            destination=t,
+                            queue_id=f"kafka:{t}",
+                            handler_canonical_id=self.current_function,
+                            file_path=file_path,
+                            line=node.lineno,
+                            handler_snippet=get_line_snippet(node.lineno),
+                            confidence="HIGH",
+                            metadata={"method": "subscribe"},
+                        )
+                    )
+
+            # Check for KafkaConsumer(topic, ...)
+            elif isinstance(node.func, ast.Name) and node.func.id == "KafkaConsumer":
+                if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
+                    t = node.args[0].value
+                    consumers.append(
+                        AsyncConsumerRecord(
+                            system=AsyncMessagingSystem.KAFKA,
+                            destination=t,
+                            queue_id=f"kafka:{t}",
+                            handler_canonical_id=self.current_function,
+                            file_path=file_path,
+                            line=node.lineno,
+                            handler_snippet=get_line_snippet(node.lineno),
+                            confidence="HIGH",
+                            metadata={"method": "KafkaConsumer"},
+                        )
+                    )
+
+            self.generic_visit(node)
+
+    KafkaCallVisitor().visit(tree)
+
+    return producers, consumers

codegraph/async_queue/detectors/rabbitmq.py

@@ -0,0 +1,116 @@
+"""RabbitMQ / AMQP / Pika message producer and consumer detector.
+
+Detects:
+  - Publishers / Producers: channel.basic_publish(exchange=..., routing_key=...)
+  - Consumers / Workers: channel.basic_consume(queue=..., on_message_callback=...)
+"""
+from __future__ import annotations
+
+import ast
+from typing import Any
+
+from codegraph.async_queue.models import (
+    AsyncConsumerRecord,
+    AsyncMessagingSystem,
+    AsyncProducerRecord,
+)
+
+
+def detect_rabbitmq(
+    tree: ast.AST,
+    content: str,
+    file_path: str,
+    module_name: str,
+) -> tuple[list[AsyncProducerRecord], list[AsyncConsumerRecord]]:
+    producers: list[AsyncProducerRecord] = []
+    consumers: list[AsyncConsumerRecord] = []
+    lines = content.splitlines()
+
+    def get_line_snippet(lineno: int) -> str:
+        if 1 <= lineno <= len(lines):
+            return lines[lineno - 1].strip()
+        return ""
+
+    class RabbitMQVisitor(ast.NodeVisitor):
+        def __init__(self) -> None:
+            self.current_function: str = f"{module_name}.<module>"
+
+        def visit_FunctionDef(self, fn: ast.FunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_AsyncFunctionDef(self, fn: ast.AsyncFunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_Call(self, node: ast.Call) -> None:
+            if isinstance(node.func, ast.Attribute):
+                attr = node.func.attr
+                # Producer: basic_publish
+                if attr in ("basic_publish", "publish"):
+                    routing_key = ""
+                    exchange = ""
+                    meta: dict[str, Any] = {"method": attr}
+                    for kw in node.keywords:
+                        if kw.arg == "routing_key" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                            routing_key = kw.value.value
+                        elif kw.arg == "exchange" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                            exchange = kw.value.value
+
+                    dest = routing_key or exchange
+                    if dest:
+                        if exchange:
+                            meta["exchange"] = exchange
+                        if routing_key:
+                            meta["routing_key"] = routing_key
+
+                        producers.append(
+                            AsyncProducerRecord(
+                                system=AsyncMessagingSystem.RABBITMQ,
+                                destination=dest,
+                                queue_id=f"rabbitmq:{dest}",
+                                caller_canonical_id=self.current_function,
+                                file_path=file_path,
+                                line=node.lineno,
+                                call_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata=meta,
+                            )
+                        )
+
+                # Consumer: basic_consume
+                elif attr in ("basic_consume", "consume"):
+                    queue_name = ""
+                    callback_fn_name = ""
+                    for kw in node.keywords:
+                        if kw.arg == "queue" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                            queue_name = kw.value.value
+                        elif kw.arg == "on_message_callback":
+                            if isinstance(kw.value, ast.Name):
+                                callback_fn_name = kw.value.id
+
+                    if queue_name:
+                        handler_id = f"{module_name}.{callback_fn_name}" if callback_fn_name else self.current_function
+                        consumers.append(
+                            AsyncConsumerRecord(
+                                system=AsyncMessagingSystem.RABBITMQ,
+                                destination=queue_name,
+                                queue_id=f"rabbitmq:{queue_name}",
+                                handler_canonical_id=handler_id,
+                                file_path=file_path,
+                                line=node.lineno,
+                                handler_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata={"callback": callback_fn_name} if callback_fn_name else {},
+                            )
+                        )
+
+            self.generic_visit(node)
+
+    RabbitMQVisitor().visit(tree)
+
+    return producers, consumers

codegraph/async_queue/detectors/redis.py

@@ -0,0 +1,130 @@
+"""Redis Pub/Sub and Streams async message detector.
+
+Detects:
+  - Publishers / Producers: r.publish(channel, message), r.xadd(stream, fields)
+  - Consumers / Workers: pubsub.subscribe(channel), r.xread({stream: id})
+"""
+from __future__ import annotations
+
+import ast
+from typing import Any
+
+from codegraph.async_queue.models import (
+    AsyncConsumerRecord,
+    AsyncMessagingSystem,
+    AsyncProducerRecord,
+)
+
+
+def detect_redis(
+    tree: ast.AST,
+    content: str,
+    file_path: str,
+    module_name: str,
+) -> tuple[list[AsyncProducerRecord], list[AsyncConsumerRecord]]:
+    producers: list[AsyncProducerRecord] = []
+    consumers: list[AsyncConsumerRecord] = []
+    lines = content.splitlines()
+
+    def get_line_snippet(lineno: int) -> str:
+        if 1 <= lineno <= len(lines):
+            return lines[lineno - 1].strip()
+        return ""
+
+    class RedisVisitor(ast.NodeVisitor):
+        def __init__(self) -> None:
+            self.current_function: str = f"{module_name}.<module>"
+
+        def visit_FunctionDef(self, fn: ast.FunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_AsyncFunctionDef(self, fn: ast.AsyncFunctionDef) -> None:
+            prev = self.current_function
+            self.current_function = f"{module_name}.{fn.name}"
+            self.generic_visit(fn)
+            self.current_function = prev
+
+        def visit_Call(self, node: ast.Call) -> None:
+            if isinstance(node.func, ast.Attribute):
+                attr = node.func.attr
+                # Producer: r.publish("channel", ...), r.xadd("stream", ...)
+                if attr in ("publish", "xadd"):
+                    channel_or_stream = ""
+                    if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str):
+                        channel_or_stream = node.args[0].value
+                    for kw in node.keywords:
+                        if kw.arg in ("channel", "name") and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                            channel_or_stream = kw.value.value
+
+                    if channel_or_stream:
+                        meta: dict[str, Any] = {"method": attr, "type": "stream" if attr == "xadd" else "pubsub"}
+                        producers.append(
+                            AsyncProducerRecord(
+                                system=AsyncMessagingSystem.REDIS,
+                                destination=channel_or_stream,
+                                queue_id=f"redis:{channel_or_stream}",
+                                caller_canonical_id=self.current_function,
+                                file_path=file_path,
+                                line=node.lineno,
+                                call_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata=meta,
+                            )
+                        )
+
+                # Consumer: pubsub.subscribe("channel"), r.xread({"stream": ...})
+                elif attr == "subscribe":
+                    channels: list[str] = []
+                    for arg in node.args:
+                        if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
+                            channels.append(arg.value)
+                        elif isinstance(arg, (ast.List, ast.Tuple)):
+                            for elt in arg.elts:
+                                if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
+                                    channels.append(elt.value)
+
+                    for ch in channels:
+                        consumers.append(
+                            AsyncConsumerRecord(
+                                system=AsyncMessagingSystem.REDIS,
+                                destination=ch,
+                                queue_id=f"redis:{ch}",
+                                handler_canonical_id=self.current_function,
+                                file_path=file_path,
+                                line=node.lineno,
+                                handler_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata={"type": "pubsub"},
+                            )
+                        )
+
+                elif attr == "xread":
+                    streams: list[str] = []
+                    if node.args and isinstance(node.args[0], ast.Dict):
+                        for k in node.args[0].keys:
+                            if isinstance(k, ast.Constant) and isinstance(k.value, str):
+                                streams.append(k.value)
+
+                    for st in streams:
+                        consumers.append(
+                            AsyncConsumerRecord(
+                                system=AsyncMessagingSystem.REDIS,
+                                destination=st,
+                                queue_id=f"redis:{st}",
+                                handler_canonical_id=self.current_function,
+                                file_path=file_path,
+                                line=node.lineno,
+                                handler_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata={"type": "stream"},
+                            )
+                        )
+
+            self.generic_visit(node)
+
+    RedisVisitor().visit(tree)
+
+    return producers, consumers