Add admin API interfaces

Author: kenliao94

src/mcp_server_rabbitmq/admin.py

@@ -0,0 +1,83 @@
+import requests
+from typing import Optional, List, Dict
+import base64
+
+class RabbitMQAdmin:
+    def __init__(self, host: str, port: int, username: str, password: str, use_tls: bool):
+        self.protocol = "https" if use_tls else "http"
+        self.base_url = f"{self.protocol}://{host}:{port}/api"
+        self.auth = base64.b64encode(f"{username}:{password}".encode()).decode()
+        self.headers = {
+            "Authorization": f"Basic {self.auth}",
+            "Content-Type": "application/json"
+        }
+
+    def _make_request(self, method: str, endpoint: str, data: Optional[Dict] = None) -> requests.Response:
+        url = f"{self.base_url}/{endpoint}"
+        response = requests.request(method, url, headers=self.headers, json=data, verify=True)
+        response.raise_for_status()
+        return response
+
+    def list_queues(self) -> List[Dict]:
+        """List all queues in the RabbitMQ server"""
+        response = self._make_request("GET", "queues")
+        return response.json()
+
+    def list_exchanges(self) -> List[Dict]:
+        """List all exchanges in the RabbitMQ server"""
+        response = self._make_request("GET", "exchanges")
+        return response.json()
+
+    def get_queue_info(self, queue: str, vhost: str = "/") -> Dict:
+        """Get detailed information about a specific queue"""
+        vhost_encoded = requests.utils.quote(vhost, safe='')
+        response = self._make_request("GET", f"queues/{vhost_encoded}/{queue}")
+        return response.json()
+
+    def delete_queue(self, queue: str, vhost: str = "/") -> None:
+        """Delete a queue"""
+        validate_rabbitmq_name(queue, "Queue name")
+        vhost_encoded = requests.utils.quote(vhost, safe='')
+        self._make_request("DELETE", f"queues/{vhost_encoded}/{queue}")
+
+    def purge_queue(self, queue: str, vhost: str = "/") -> None:
+        """Remove all messages from a queue"""
+        validate_rabbitmq_name(queue, "Queue name")
+        vhost_encoded = requests.utils.quote(vhost, safe='')
+        self._make_request("DELETE", f"queues/{vhost_encoded}/{queue}/contents")
+
+    def get_exchange_info(self, exchange: str, vhost: str = "/") -> Dict:
+        """Get detailed information about a specific exchange"""
+        vhost_encoded = requests.utils.quote(vhost, safe='')
+        response = self._make_request("GET", f"exchanges/{vhost_encoded}/{exchange}")
+        return response.json()
+
+    def delete_exchange(self, exchange: str, vhost: str = "/") -> None:
+        """Delete an exchange"""
+        validate_rabbitmq_name(exchange, "Exchange name")
+        vhost_encoded = requests.utils.quote(vhost, safe='')
+        self._make_request("DELETE", f"exchanges/{vhost_encoded}/{exchange}")
+
+    def get_bindings(self, queue: Optional[str] = None, exchange: Optional[str] = None, vhost: str = "/") -> List[Dict]:
+        """Get bindings, optionally filtered by queue or exchange"""
+        vhost_encoded = requests.utils.quote(vhost, safe='')
+        if queue:
+            validate_rabbitmq_name(queue, "Queue name")
+            response = self._make_request("GET", f"queues/{vhost_encoded}/{queue}/bindings")
+        elif exchange:
+            validate_rabbitmq_name(exchange, "Exchange name")
+            response = self._make_request("GET", f"exchanges/{vhost_encoded}/{exchange}/bindings/source")
+        else:
+            response = self._make_request("GET", f"bindings/{vhost_encoded}")
+        return response.json()
+
+    def get_overview(self) -> Dict:
+        """Get overview of RabbitMQ server including version, stats, and listeners"""
+        response = self._make_request("GET", "overview")
+        return response.json()
+    
+if __name__ == "__main__":
+    # test
+    admin = RabbitMQAdmin("localhost", 15672, "guest", "guest", False)
+    print(admin.list_queues())
+    print("Done")

src/mcp_server_rabbitmq/handlers.py

@@ -0,0 +1,13 @@
+from .connection import RabbitMQConnection
+
+def handle_enqueue(rabbitmq: RabbitMQConnection, queue: str, message: str):
+    connection, channel = rabbitmq.get_channel()
+    channel.queue_declare(queue)
+    channel.basic_publish(exchange="", routing_key=queue, body=message)
+    connection.close()
+
+def handle_fanout(rabbitmq: RabbitMQConnection, exchange: str, message: str):
+    connection, channel = rabbitmq.get_channel()
+    channel.exchange_declare(exchange=exchange, exchange_type="fanout")
+    channel.basic_publish(exchange=exchange, routing_key="", body=message)
+    connection.close()

src/mcp_server_rabbitmq/server.py

@@ -9,6 +9,7 @@
 from .models import Enqueue, Fanout
 from .logger import Logger, LOG_LEVEL
 from .connection import RabbitMQConnection, validate_rabbitmq_name
+from .handlers import handle_enqueue, handle_fanout
 
 
 async def serve(rabbitmq_host: str, port: int, username: str, password: str, use_tls: bool, log_level: str = LOG_LEVEL.DEBUG.name) -> None:
@@ -55,10 +56,7 @@ async def call_tool(
             validate_rabbitmq_name(queue, "Queue name")
 
             try:
-                connection, channel = rabbitmq.get_channel()
-                channel.queue_declare(queue)
-                channel.basic_publish(exchange="", routing_key=queue, body=message)
-                connection.close()
+                handle_enqueue(rabbitmq, queue, message)
                 return [TextContent(type="text", text=str("suceeded"))]
             except Exception as e:
                 logger.error(f"{e}")
@@ -71,10 +69,7 @@ async def call_tool(
             validate_rabbitmq_name(exchange, "Exchange name")
 
             try:
-                connection, channel = rabbitmq.get_channel()
-                channel.exchange_declare(exchange=exchange, exchange_type="fanout")
-                channel.basic_publish(exchange=exchange, routing_key="", body=message)
-                connection.close()
+                handle_fanout(rabbitmq, exchange, message)
                 return [TextContent(type="text", text=str("suceeded"))]
             except Exception as e:
                 logger.error(f"{e}")