Add log level as configuration

Author: kenliao94

src/mcp_server_rabbitmq/__init__.py

@@ -13,14 +13,16 @@ def main():
     parser.add_argument("--username", type=str, help="Username for the connection")
     parser.add_argument("--password", type=str, help="Password for the connection")
     parser.add_argument("--use-tls", type=bool, help="Is the connection amqps")
+    parser.add_argument("--log-level", type=str, help="Log level, supports DEBUG|INFO|WARNING|ERROR, default to WARNING")
 
     args = parser.parse_args()
     asyncio.run(serve(
         rabbitmq_host=args.rabbitmq_host,
         port=args.port,
         username=args.username,
         password=args.password,
-        use_tls=args.use_tls))
+        use_tls=args.use_tls,
+        log_level=args.log_level))
 
 
 if __name__ == "__main__":

src/mcp_server_rabbitmq/logger.py

@@ -31,7 +31,7 @@ def error(self, msg):
 
     def __post_process_log_entry(self, msg: str, log_level: LOG_LEVEL) -> str:
         timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        return f"[{log_level}] {timestamp} - {msg}\n"
+        return f"[{log_level.name}] {timestamp} - {msg}\n"
 
     def __write_log(self, msg):
         with open(self.log_file_path, "a") as f:

src/mcp_server_rabbitmq/server.py

@@ -10,11 +10,19 @@
 from .logger import Logger, LOG_LEVEL
 
 
-async def serve(rabbitmq_host: str, port: int, username: str, password: str, use_tls: bool, logging_level: LOG_LEVEL = LOG_LEVEL.WARNING) -> None:
-    # Setup server and logger
+async def serve(rabbitmq_host: str, port: int, username: str, password: str, use_tls: bool, log_level: str = LOG_LEVEL.DEBUG.name) -> None:
+    # Setup server
     server = Server("mcp-rabbitmq")
-    # Make logger another option to pass
-    logger = Logger("server.log")
+    # Setup logger
+    is_log_level_exception = False
+    try:
+        log_level = LOG_LEVEL[log_level]
+    except Exception:
+        is_log_level_exception = True
+        log_level = LOG_LEVEL.WARNING
+    logger = Logger("server.log", log_level)
+    if is_log_level_exception:
+        logger.warning("Wrong log_level received. Default to WARNING")
     # Setup RabbitMQ connection metadata
     protocol = "amqps" if use_tls else "amqp"
     url = f"{protocol}://{username}:{password}@{rabbitmq_host}:{port}"
@@ -40,6 +48,7 @@ async def call_tool(
         arguments: dict
     ) -> list[TextContent]:
         if name == "enqueue":
+            logger.debug("Executing enqueue tool")
             message = arguments["message"]
             queue = arguments["queue"]
             # Send to RabbitMQ host
@@ -50,7 +59,7 @@ async def call_tool(
                 channel.basic_publish(exchange="", routing_key=queue, body=message)
                 return [TextContent(type="text", text=str("suceeded"))]
             except Exception as e:
-                logger.error(f"[ERROR] {e}")
+                logger.error(f"{e}")
                 return [TextContent(type="text", text=str("failed"))]
         raise ValueError(f"Tool not found: {name}")