diff --git a/README.md b/README.md index 60b1035..5d8d94d 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,24 @@ The MCP server supports multiple transport mechanisms. For detailed information | `sse` (legacy) | Server-Sent Events | Streaming applications | | `streamable-http` | Streamable HTTP transport | Container deployments, remote servers | +## Health Endpoint + +The MCP server provides a `/health` endpoint that returns a simple 200 OK response with a JSON body indicating the server's status. This endpoint is available when using HTTP-based transports (`http`, `sse`, `streamable-http`) and is useful for: + +- AWS Application Load Balancer (ALB) health checks +- Kubernetes liveness and readiness probes +- Docker health checks +- Other orchestration systems + +**Example Response:** +```json +{ + "status": "ok" +} +``` + +This endpoint does not require authentication and can be accessed directly via HTTP GET: + ## Usage ```bash @@ -338,8 +356,11 @@ docker ps # Check container logs docker logs mcp-server-snowflake -# Test endpoint (should return MCP server info) +# Test MCP endpoint (should return MCP server info) curl http://localhost:9000/snowflake-mcp + +# Test health endpoint (should return {"status": "ok"}) +curl http://localhost:9000/health ``` ## Docker Compose Deployment @@ -387,8 +408,11 @@ docker-compose ps # View logs docker-compose logs -# Test endpoint +# Test MCP endpoint curl http://localhost:9000/snowflake-mcp + +# Test health endpoint +curl http://localhost:9000/health ``` ## Connecting MCP Clients to Containers diff --git a/mcp_server_snowflake/server.py b/mcp_server_snowflake/server.py index e3053f3..a09ad56 100644 --- a/mcp_server_snowflake/server.py +++ b/mcp_server_snowflake/server.py @@ -22,6 +22,8 @@ from fastmcp.utilities.logging import get_logger from snowflake.connector import DictCursor, connect from snowflake.core import Root +from starlette.requests import Request +from starlette.responses import JSONResponse from mcp_server_snowflake.cortex_services.tools import ( initialize_cortex_agent_tool, @@ -581,6 +583,19 @@ def initialize_tools(snowflake_service: SnowflakeService, server: FastMCP): initialize_cortex_analyst_tool(server, snowflake_service) +def add_health_endpoint(server: FastMCP): + """ + Add a health endpoint to the server for compatibility with load balancers and orchestration systems. + + This endpoint returns a simple 200 OK response with a JSON body indicating the server's status. + It can be used by AWS ALB, Kubernetes, or other systems that expect a standard HTTP health check. + """ + + @server.custom_route("/health", methods=["GET"]) + async def health(request: Request) -> JSONResponse: + return JSONResponse({"status": "ok"}) + + def main(): args = parse_arguments() @@ -598,6 +613,8 @@ def main(): "streamable-http", ]: endpoint = os.environ.get("SNOWFLAKE_MCP_ENDPOINT", args.endpoint) + # Add health endpoint for load balancers and orchestration systems + add_health_endpoint(server) logger.info(f"Starting server with transport: {args.transport}") server.run( transport=args.transport, host="0.0.0.0", port=9000, path=endpoint diff --git a/mcp_server_snowflake/tests/test_health_endpoint.py b/mcp_server_snowflake/tests/test_health_endpoint.py new file mode 100644 index 0000000..8a9fa58 --- /dev/null +++ b/mcp_server_snowflake/tests/test_health_endpoint.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python +# Copyright 2025 Snowflake Inc. +# SPDX-License-Identifier: Apache-2.0 +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" +Test for the health endpoint of the MCP server. +""" + +import os +import signal +import subprocess +import sys +import time +import unittest +from pathlib import Path + +import requests + + +class TestHealthEndpoint(unittest.TestCase): + """Test the health endpoint of the MCP server.""" + + @classmethod + def setUpClass(cls): + """Start the MCP server before running tests.""" + # Path to the configuration file + config_file = ( + Path(__file__).parent.parent.parent / "services" / "configuration.yaml" + ) + + # Start the server with HTTP transport + cls.server_process = subprocess.Popen( + [ + sys.executable, + "-m", + "mcp_server_snowflake.server", + "--transport", + "http", + "--service-config-file", + str(config_file), + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + # Wait for the server to start + time.sleep(2) + + @classmethod + def tearDownClass(cls): + """Stop the MCP server after running tests.""" + if cls.server_process: + os.kill(cls.server_process.pid, signal.SIGTERM) + cls.server_process.wait() + + def test_health_endpoint(self): + """Test that the health endpoint returns a 200 OK response.""" + response = requests.get("http://localhost:9000/health") + self.assertEqual(response.status_code, 200) + self.assertEqual(response.json(), {"status": "ok"}) + + +if __name__ == "__main__": + unittest.main()