codegraph-engine 3.0.0__2-py3-none-any.whl

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

codegraph/async_queue/detectors/sqs.py

@@ -0,0 +1,135 @@
+"""AWS SQS message producer and consumer detector.
+
+Detects:
+  - Publishers / Producers: sqs.send_message(QueueUrl=...), sqs.send_message_batch()
+  - Consumers / Workers: sqs.receive_message(QueueUrl=...), @sqs_handler()
+"""
+from __future__ import annotations
+
+import ast
+from typing import Any
+
+from codegraph.async_queue.models import (
+    AsyncConsumerRecord,
+    AsyncMessagingSystem,
+    AsyncProducerRecord,
+)
+
+
+def _extract_queue_name(url_or_name: str) -> str:
+    """Extract queue name from SQS URL or name string."""
+    clean = url_or_name.strip().rstrip("/")
+    if "/" in clean:
+        return clean.split("/")[-1]
+    return clean
+
+
+def detect_sqs(
+    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 SQSVisitor(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}"
+            # Check for decorator like @sqs_handler
+            for dec in fn.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 "sqs" in dec_name.lower():
+                    q_name = ""
+                    if isinstance(dec, ast.Call):
+                        for kw in dec.keywords:
+                            if kw.arg in ("queue", "queue_name", "queue_url") and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                                q_name = _extract_queue_name(kw.value.value)
+                    if q_name:
+                        consumers.append(
+                            AsyncConsumerRecord(
+                                system=AsyncMessagingSystem.AWS_SQS,
+                                destination=q_name,
+                                queue_id=f"aws_sqs:{q_name}",
+                                handler_canonical_id=f"{module_name}.{fn.name}",
+                                file_path=file_path,
+                                line=fn.lineno,
+                                handler_snippet=get_line_snippet(fn.lineno),
+                                confidence="HIGH",
+                                metadata={"decorator": dec_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: send_message, send_message_batch
+                if attr in ("send_message", "send_message_batch"):
+                    q_name = ""
+                    meta: dict[str, Any] = {"method": attr}
+                    for kw in node.keywords:
+                        if kw.arg == "QueueUrl" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                            q_name = _extract_queue_name(kw.value.value)
+                        elif kw.arg in ("DelaySeconds", "MessageDeduplicationId", "MessageGroupId"):
+                            meta[kw.arg] = True
+
+                    if q_name:
+                        producers.append(
+                            AsyncProducerRecord(
+                                system=AsyncMessagingSystem.AWS_SQS,
+                                destination=q_name,
+                                queue_id=f"aws_sqs:{q_name}",
+                                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: receive_message
+                elif attr == "receive_message":
+                    q_name = ""
+                    for kw in node.keywords:
+                        if kw.arg == "QueueUrl" and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str):
+                            q_name = _extract_queue_name(kw.value.value)
+                    if q_name:
+                        consumers.append(
+                            AsyncConsumerRecord(
+                                system=AsyncMessagingSystem.AWS_SQS,
+                                destination=q_name,
+                                queue_id=f"aws_sqs:{q_name}",
+                                handler_canonical_id=self.current_function,
+                                file_path=file_path,
+                                line=node.lineno,
+                                handler_snippet=get_line_snippet(node.lineno),
+                                confidence="HIGH",
+                                metadata={"method": "receive_message"},
+                            )
+                        )
+
+            self.generic_visit(node)
+
+    SQSVisitor().visit(tree)
+
+    return producers, consumers