diff --git a/src/server/databricks_mcp_server.py b/src/server/databricks_mcp_server.py index 4534c97..a1fb155 100644 --- a/src/server/databricks_mcp_server.py +++ b/src/server/databricks_mcp_server.py @@ -34,9 +34,7 @@ class DatabricksMCPServer(FastMCP): def __init__(self): """Initialize the Databricks MCP server.""" - super().__init__(name="databricks-mcp", - version="1.0.0", - instructions="Use this server to manage Databricks resources") + super().__init__(name="databricks-mcp") logger.info("Initializing Databricks MCP server") logger.info(f"Databricks host: {settings.DATABRICKS_HOST}") @@ -55,10 +53,10 @@ async def list_clusters(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Listing clusters with params: {params}") try: result = await clusters.list_clusters() - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error listing clusters: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] @self.tool( name="create_cluster", @@ -68,10 +66,10 @@ async def create_cluster(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Creating cluster with params: {params}") try: result = await clusters.create_cluster(params) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error creating cluster: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] @self.tool( name="terminate_cluster", @@ -81,10 +79,10 @@ async def terminate_cluster(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Terminating cluster with params: {params}") try: result = await clusters.terminate_cluster(params.get("cluster_id")) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error terminating cluster: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] @self.tool( name="get_cluster", @@ -94,10 +92,10 @@ async def get_cluster(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Getting cluster info with params: {params}") try: result = await clusters.get_cluster(params.get("cluster_id")) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error getting cluster info: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] @self.tool( name="start_cluster", @@ -107,10 +105,10 @@ async def start_cluster(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Starting cluster with params: {params}") try: result = await clusters.start_cluster(params.get("cluster_id")) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error starting cluster: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] # Job management tools @self.tool( @@ -121,10 +119,10 @@ async def list_jobs(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Listing jobs with params: {params}") try: result = await jobs.list_jobs() - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error listing jobs: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] @self.tool( name="run_job", @@ -135,10 +133,10 @@ async def run_job(params: Dict[str, Any]) -> List[TextContent]: try: notebook_params = params.get("notebook_params", {}) result = await jobs.run_job(params.get("job_id"), notebook_params) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error running job: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] # Notebook management tools @self.tool( @@ -149,10 +147,10 @@ async def list_notebooks(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Listing notebooks with params: {params}") try: result = await notebooks.list_notebooks(params.get("path")) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error listing notebooks: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] @self.tool( name="export_notebook", @@ -170,10 +168,10 @@ async def export_notebook(params: Dict[str, Any]) -> List[TextContent]: summary = f"{content[:1000]}... [content truncated, total length: {len(content)} characters]" result["content"] = summary - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error exporting notebook: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] # DBFS tools @self.tool( @@ -184,10 +182,10 @@ async def list_files(params: Dict[str, Any]) -> List[TextContent]: logger.info(f"Listing files with params: {params}") try: result = await dbfs.list_files(params.get("dbfs_path")) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error listing files: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] # SQL tools @self.tool( @@ -203,10 +201,10 @@ async def execute_sql(params: Dict[str, Any]) -> List[TextContent]: schema = params.get("schema") result = await sql.execute_sql(statement, warehouse_id, catalog, schema) - return [{"text": json.dumps(result)}] + return [{"type": "text", "text": json.dumps(result)}] except Exception as e: logger.error(f"Error executing SQL: {str(e)}") - return [{"text": json.dumps({"error": str(e)})}] + return [{"type": "text", "text": json.dumps({"error": str(e)})}] async def main():