Snowflake Cube Server

from __future__ import annotations from typing import Annotated, Any, Literal, Optional, Union from mcp.server.fastmcp import FastMCP from mcp.types import TextContent, EmbeddedResource, TextResourceContents import jwt from pydantic import BaseModel, Field import requests import json import logging import yaml import uuid import time def data_to_yaml(data: Any) -> str: return yaml.dump(data, indent=2, sort_keys=False) class CubeClient: Route = Literal["meta", "load"] max_wait_time = 10 request_backoff = 1 def __init__(self, endpoint: str, api_secret: str, token_payload: dict, logger: logging.Logger): self.endpoint = endpoint self.api_secret = api_secret self.token_payload = token_payload self.token = None self.logger = logger self._refresh_token() self.meta = self.describe() def _generate_token(self): return jwt.encode(self.token_payload, self.api_secret, algorithm="HS256") def _refresh_token(self): self.token = self._generate_token() def _request(self, route: Route, **params): request_time = time.time() headers = {"Authorization": self.token} url = f"{self.endpoint if self.endpoint[-1] != '/' else self.endpoint[:-1]}/{route}" serialized_params = {k: json.dumps(v) for k, v in params.items()} try: response = requests.get(url, headers=headers, params=serialized_params) # Handle "continue wait" responses while response.json().get("error") == "Continue wait": if time.time() - request_time > self.max_wait_time: self.logger.error(f"Request timed out after {self.max_wait_time} seconds") return {"error": "Request timed out. Something may have gone wrong or the request may be too complex."} self.logger.warning(f"Request incomplete, polling again in {self.request_backoff} second(s)") time.sleep(self.request_backoff) response = requests.get(url, headers=headers, params=serialized_params) # Handle 403 responses by trying to refresh the token once if response.status_code == 403: self.logger.warning("Received 403, attempting token refresh") self._refresh_token() return requests.get(url, headers=headers, params=serialized_params) if response.status_code != 200: self.logger.error(f"Request failed with error: {str(response.json().get('error'))}") return response.json() except Exception as e: self.logger.error(f"Request failed with error: {str(e)}") return {"error": f"Request failed: {str(e)}"} def describe(self): return self._request("meta") def _cast_numerics(self, response): if response.get("data") and response.get("annotation"): # Find which keys are numeric numeric_keys = set() dimensions_and_measures = dict( response["annotation"].get("dimensions", {}), **response["annotation"].get("measures", []) ) for column_name, column in dimensions_and_measures.items(): if column.get("type") == "number": numeric_keys.add(column_name) # Cast numeric values to numbers for row in response["data"]: for key in numeric_keys: try: row[key] = float(row[key]) if row[key].is_integer(): row[key] = int(row[key]) except (ValueError, TypeError): pass return response def query(self, query, cast_numerics=True): response = self._request("load", query=query) if cast_numerics: response = self._cast_numerics(response) return response class Filter(BaseModel): dimension: str = Field(..., description="Name of the time dimension") granularity: Literal["second", "minute", "hour", "day", "week", "month", "quarter", "year"] = Field( ..., description="Time granularity" ) dateRange: Union[list[str], str] = Field( ..., description="Pair of dates ISO dates representing the start and end of the range. Alternatively, a string representing a relative date range of the form: 'last N days', 'today', 'yesterday', 'last year', etc.", ) model_config = {"exclude_none": True} class TimeDimension(BaseModel): dimension: str = Field(..., description="Name of the time dimension") granularity: Literal["second", "minute", "hour", "day", "week", "month", "quarter", "year"] = Field( ..., description="Time granularity" ) dateRange: Union[list[str], str] = Field( ..., description="Pair of dates ISO dates representing the start and end of the range. Alternatively, a string representing a relative date range of the form: 'last N days', 'today', 'yesterday', 'last year', etc.", ) model_config = {"exclude_none": True} class Query(BaseModel): measures: list[str] = Field([], description="Names of measures to query") dimensions: list[str] = Field([], description="Names of dimensions to group by") timeDimensions: list[TimeDimension] = Field([], description="Time dimensions to group by") # filters: list[Filter] = Field([], description="Filters to apply to the query") limit: Optional[int] = Field(500, description="Maximum number of rows to return. Defaults to 500") offset: Optional[int] = Field(0, description="Number of rows to skip. Defaults to 0") order: dict[str, Literal["asc", "desc"]] = Field( {}, description="Optional ordering of the results. The order is sensitive to the order of keys." ) ungrouped: bool = Field( False, description="Return results without grouping by dimensions. Instead, return all rows. This can be useful for fetching a single row by its ID as well.", ) model_config = {"exclude_none": True} def main(credentials, logger): mcp = FastMCP("Cube.dev") client = CubeClient(**credentials, logger=logger) @mcp.resource("context://data_description") def data_description() -> str: """Describe the data available in Cube.""" meta = client.describe() if error := meta.get("error"): logger.error("Error in data_description: %s\n\n%s", error, meta.get("stack")) logger.error("Full response: %s", json.dumps(meta)) return f"Error: Description of the data is not available: {error}, {meta}" description = [ { "name": cube.get("name"), "title": cube.get("title"), "description": cube.get("description"), "dimensions": [ { "name": dimension.get("name"), "title": dimension.get("shortTitle") or dimension.get("title"), "description": dimension.get("description"), } for dimension in cube.get("dimensions", []) ], "measures": [ { "name": measure.get("name"), "title": measure.get("shortTitle") or measure.get("title"), "description": measure.get("description"), } for measure in cube.get("measures", []) ], } for cube in meta.get("cubes", []) ] return "Here is a description of the data available via the read_data tool:\n\n" + yaml.dump( description, indent=2, sort_keys=True ) @mcp.tool("describe_data") def describe_data() -> str: """Describe the data available in Cube.""" return {"type": "text", "text": data_description()} @mcp.tool("read_data") def read_data(query: Query) -> str: """Read data from Cube.""" try: query_dict = query.model_dump(by_alias=True, exclude_none=True) logger.info("read_data called with query: %s", json.dumps(query_dict)) response = client.query(query_dict) if error := response.get("error"): logger.error("Error in read_data: %s\n\n%s", error, response.get("stack")) logger.error("Full response: %s", json.dumps(response)) return f"Error: {error}" data = response.get("data", []) logger.info("read_data returned %s rows", len(data)) data_id = str(uuid.uuid4()) @mcp.resource(f"data://{data_id}") def data_resource() -> str: return json.dumps(data) logger.info("Added results as resource with ID: %s", data_id) output = { "type": "data", "data_id": data_id, "data": data, } yaml_output = data_to_yaml(output) json_output = json.dumps(output) return [ TextContent(type="text", text=yaml_output), EmbeddedResource( type="resource", resource=TextResourceContents(uri=f"data://{data_id}", text=json_output, mimeType="application/json"), ), ] except Exception as e: logger.error("Error in read_data: %s", str(e)) return f"Error: {str(e)}" logger.info("Starting Cube MCP server") mcp.run()