From b4af8ea32c752760c43b22179b1c22afdb8368f8 Mon Sep 17 00:00:00 2001 From: Sean Redmond Date: Sun, 26 Oct 2025 13:10:54 +0000 Subject: [PATCH 1/3] fix: add health endpoint for compatibility with load balancers This commit adds a endpoint to the MCP server that returns a 200 OK response with a JSON payload {"status": "ok"}. This endpoint enables compatibility with AWS ALB, Kubernetes, and other orchestration systems that require standard HTTP health checks. Fixes #138 --- README.md | 28 +++++++- mcp_server_snowflake/server.py | 15 ++++ .../tests/test_health_endpoint.py | 68 +++++++++++++++++++ 3 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 mcp_server_snowflake/tests/test_health_endpoint.py 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..f93bbff 100644 --- a/mcp_server_snowflake/server.py +++ b/mcp_server_snowflake/server.py @@ -581,6 +581,18 @@ 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.app.get("/health") + async def health(): + return {"status": "ok"} + + def main(): args = parse_arguments() @@ -588,6 +600,9 @@ def main(): # Create server with lifespan that has access to args server = FastMCP("Snowflake MCP Server", lifespan=create_lifespan(args)) + + # Add health endpoint for load balancers and orchestration systems + add_health_endpoint(server) try: logger.info("Starting Snowflake MCP Server...") 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..58a6595 --- /dev/null +++ b/mcp_server_snowflake/tests/test_health_endpoint.py @@ -0,0 +1,68 @@ +#!/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 unittest +import requests +import subprocess +import time +import os +import signal +import sys +from pathlib import Path + + +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() From 3371dc4c9213d9dc52b59fb3b1d400ee21b76a75 Mon Sep 17 00:00:00 2001 From: Jason Summer Date: Thu, 30 Oct 2025 10:51:18 -0500 Subject: [PATCH 2/3] fix: Use custom_route() for health endpoint --- mcp_server_snowflake/server.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/mcp_server_snowflake/server.py b/mcp_server_snowflake/server.py index f93bbff..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, @@ -584,13 +586,14 @@ def initialize_tools(snowflake_service: SnowflakeService, server: FastMCP): 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.app.get("/health") - async def health(): - return {"status": "ok"} + + @server.custom_route("/health", methods=["GET"]) + async def health(request: Request) -> JSONResponse: + return JSONResponse({"status": "ok"}) def main(): @@ -600,9 +603,6 @@ def main(): # Create server with lifespan that has access to args server = FastMCP("Snowflake MCP Server", lifespan=create_lifespan(args)) - - # Add health endpoint for load balancers and orchestration systems - add_health_endpoint(server) try: logger.info("Starting Snowflake MCP Server...") @@ -613,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 From 52cd61d0f4c79b46437299b96f7cd61fc7c03437 Mon Sep 17 00:00:00 2001 From: Jason Summer Date: Thu, 30 Oct 2025 11:08:07 -0500 Subject: [PATCH 3/3] chore: Run ruff --- .../tests/test_health_endpoint.py | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/mcp_server_snowflake/tests/test_health_endpoint.py b/mcp_server_snowflake/tests/test_health_endpoint.py index 58a6595..8a9fa58 100644 --- a/mcp_server_snowflake/tests/test_health_endpoint.py +++ b/mcp_server_snowflake/tests/test_health_endpoint.py @@ -13,15 +13,17 @@ """ Test for the health endpoint of the MCP server. """ -import unittest -import requests -import subprocess -import time + 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.""" @@ -30,8 +32,10 @@ class TestHealthEndpoint(unittest.TestCase): 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" - + config_file = ( + Path(__file__).parent.parent.parent / "services" / "configuration.yaml" + ) + # Start the server with HTTP transport cls.server_process = subprocess.Popen( [ @@ -46,17 +50,17 @@ def setUpClass(cls): 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")