"""
FastMCP server for Prometheus MCP.
This module provides the MCP server that exposes PromQL query tools
for AWS Managed Prometheus with SigV4 authentication.
Uses FastMCP with stateless HTTP for Kubernetes deployment.
"""
import logging
import os
from typing import Any
from fastmcp import FastMCP
from starlette.requests import Request
from starlette.responses import JSONResponse
from prometheus_mcp.promql.tools import (
init_client,
get_label_values as promql_get_label_values,
list_labels as promql_list_labels,
list_metrics as promql_list_metrics,
query_instant as promql_query_instant,
query_range as promql_query_range,
)
# Configure logging
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("prometheus-mcp")
# Initialize FastMCP with stateless HTTP
mcp = FastMCP("prometheus-mcp", stateless_http=True)
@mcp.custom_route("/health", methods=["GET"])
async def health_check(request: Request) -> JSONResponse:
"""Health check endpoint for Kubernetes probes."""
return JSONResponse({"status": "healthy"})
@mcp.custom_route("/ready", methods=["GET"])
async def readiness_check(request: Request) -> JSONResponse:
"""Readiness check endpoint for Kubernetes probes."""
return JSONResponse({"status": "ready"})
def initialize_amp_client() -> None:
"""Initialize the AMP client on startup."""
workspace_id = os.environ.get("PROMETHEUS_WORKSPACE_ID")
region = os.environ.get("AWS_REGION", "us-east-1")
if not workspace_id:
logger.warning(
"PROMETHEUS_WORKSPACE_ID not set. Set it before making queries."
)
return
init_client(workspace_id=workspace_id, region=region)
logger.info(f"Initialized AMP client for workspace: {workspace_id} in {region}")
@mcp.tool()
async def query_instant(query: str, time: str | None = None) -> dict[str, Any]:
"""
Execute an instant PromQL query against AWS Managed Prometheus.
Returns current values of matching time series at a single point in time.
Use for current metric values or point-in-time snapshots.
Args:
query: PromQL query expression (e.g., 'up', 'rate(http_requests_total[5m])')
time: Optional evaluation timestamp in RFC3339 format
(e.g., '2024-01-15T10:30:00Z') or Unix timestamp.
Defaults to current server time.
Returns:
Query result containing status and matching time series with their current values.
"""
result = await promql_query_instant(query=query, time=time)
return result.model_dump()
@mcp.tool()
async def query_range(
query: str,
start: str,
end: str,
step: str = "1m",
) -> dict[str, Any]:
"""
Execute a range PromQL query to get time series data over a time period.
Returns data points at regular intervals.
Use for historical data analysis, graphing, and trend analysis.
Args:
query: PromQL query expression (e.g., 'up', 'rate(http_requests_total[5m])')
start: Start timestamp in RFC3339 format (e.g., '2024-01-15T00:00:00Z')
or Unix timestamp.
end: End timestamp in RFC3339 format (e.g., '2024-01-15T12:00:00Z')
or Unix timestamp.
step: Query resolution step width (e.g., '15s', '1m', '5m', '1h').
Defaults to '1m'. Smaller steps = more data points but slower queries.
Returns:
Range query result containing time series with values at each step interval.
"""
result = await promql_query_range(query=query, start=start, end=end, step=step)
return result.model_dump()
@mcp.tool()
async def list_labels(match: list[str] | None = None) -> dict[str, Any]:
"""
Get all label names from AWS Managed Prometheus.
Use for discovering what labels are available for filtering queries.
Args:
match: Optional list of series selectors to filter which series to consider.
For example: ['up', 'http_requests_total']
Returns:
List of all unique label names.
"""
result = await promql_list_labels(match=match)
return result.model_dump()
@mcp.tool()
async def get_label_values(
label_name: str,
match: list[str] | None = None,
) -> dict[str, Any]:
"""
Get all values for a specific label from AWS Managed Prometheus.
Use to discover possible filter values for a specific label.
Args:
label_name: The label name to get values for (e.g., 'job', 'instance', 'namespace')
match: Optional list of series selectors to filter which series to consider.
Returns:
List of all unique values for the specified label.
"""
result = await promql_get_label_values(label_name=label_name, match=match)
return result.model_dump()
@mcp.tool()
async def list_metrics(with_metadata: bool = False) -> dict[str, Any]:
"""
Get all metric names from AWS Managed Prometheus.
Optionally includes metadata like metric type, help text, and unit.
Args:
with_metadata: If True, fetches full metadata (type, help, unit) for each metric.
This is slower but provides more information. Defaults to False.
Returns:
List of all available metrics with optional metadata.
"""
result = await promql_list_metrics(with_metadata=with_metadata)
return result.model_dump()
def main() -> None:
"""Main entry point."""
initialize_amp_client()
host = os.environ.get("MCP_HOST", "0.0.0.0")
port = int(os.environ.get("MCP_PORT", "8080"))
logger.info(f"Starting FastMCP server on {host}:{port}...")
mcp.run(
transport="http",
host=host,
port=port,
log_level="info",
)
if __name__ == "__main__":
main()
