|
import os |
|
import sys |
|
import asyncio |
|
from typing import AsyncIterator |
|
from contextlib import asynccontextmanager |
|
from arango import ArangoClient |
|
from arango.exceptions import ArangoError |
|
from mcp.server.fastmcp import FastMCP |
|
from mcp.types import TextContent |
|
from loguru import logger |
|
|
|
# Configure logging |
|
logger.remove() |
|
logger.add(sys.stderr, level="INFO", format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>") |
|
|
|
# ArangoDB connection settings |
|
ARANGO_URL = os.getenv("ARANGO_URL", "http://localhost:8529") |
|
ARANGO_DB = os.getenv("ARANGO_DB", "verifaix") |
|
ARANGO_USERNAME = os.getenv("ARANGO_USERNAME", "root") |
|
ARANGO_PASSWORD = os.getenv("ARANGO_PASSWORD", "openSesame") |
|
|
|
class ArangoDBContext: |
|
"""Context manager for handling ArangoDB connection.""" |
|
def __init__(self): |
|
self.client = None |
|
self.db = None |
|
|
|
async def connect(self): |
|
"""Establish connection to ArangoDB.""" |
|
logger.info(f"π Connecting to ArangoDB at {ARANGO_URL}...") |
|
try: |
|
self.client = ArangoClient(hosts=ARANGO_URL) |
|
self.db = self.client.db(ARANGO_DB, username=ARANGO_USERNAME, password=ARANGO_PASSWORD) |
|
version = self.db.version() |
|
logger.info(f"β
Connected to ArangoDB. Version: {version}") |
|
except ArangoError as e: |
|
logger.error(f"β Failed to connect to ArangoDB: {str(e)}") |
|
raise e |
|
|
|
async def disconnect(self): |
|
"""Disconnect from ArangoDB (if needed).""" |
|
logger.info("π Closing ArangoDB connection...") |
|
self.client = None # ArangoDB client does not need explicit close |
|
logger.info("β
ArangoDB connection closed.") |
|
|
|
@asynccontextmanager |
|
async def app_lifespan(server: FastMCP) -> AsyncIterator[dict]: |
|
"""Manage application lifecycle with ArangoDB.""" |
|
context = ArangoDBContext() |
|
try: |
|
await context.connect() # β
Establish DB connection on startup |
|
yield {"db": context.db} # β
Store the DB in the lifespan context |
|
finally: |
|
await context.disconnect() # β
Cleanup on shutdown |
|
|
|
# Initialize MCP with lifespan support |
|
mcp = FastMCP( |
|
name="ArangoDB Tools", |
|
version="0.1.0", |
|
description="ArangoDB MCP Server", |
|
lifespan=app_lifespan # β
Ensure server stays alive |
|
) |
|
|
|
@mcp.tool() |
|
def list_collections(ctx) -> TextContent: |
|
"""List all collections in the ArangoDB database.""" |
|
db = ctx.request_context.lifespan_context["db"] # β
Get DB from lifespan |
|
try: |
|
collections = db.collections() |
|
collection_list = [col["name"] for col in collections] |
|
logger.info(f"π Listed {len(collection_list)} collections.") |
|
return TextContent(type="text", text=str(collection_list)) |
|
except ArangoError as e: |
|
logger.error(f"β Error listing collections: {str(e)}") |
|
return TextContent(type="text", text=f"Error: {str(e)}") |
|
|
|
if __name__ == "__main__": |
|
try: |
|
logger.info("π Starting ArangoDB MCP server...") |
|
mcp.run() # β
This should now block and keep the server running |
|
except Exception as e: |
|
logger.error(f"β MCP Server startup failed: {str(e)}") |
|
sys.exit(1) |