import logging
from dotenv import load_dotenv
from mcp.server.fastmcp import FastMCP
import os
from typing import Dict, Any
import bauplan
import datetime
env_path = os.path.join(os.getcwd(), ".env")
load_dotenv(env_path)
from mcp_bauplan.mcp_config import config
MCP_SERVER_NAME = "mcp-bauplan"
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger(MCP_SERVER_NAME)
deps = ["starlette", "python-dotenv", "uvicorn", "httpx", "bauplan"]
mcp = FastMCP(MCP_SERVER_NAME, dependencies=deps)
def create_bauplan_client():
"""
Creates and validates a connection Bauplan.
Retrieves connection parameters from config, establishes a connection.
Returns:
Client: A configured Bauplan client instance
Raises:
ConnectionError: When connection cannot be established
ConfigurationError: When configuration is invalid
"""
logger.info(
f"Creating Bauplan client connection. "
f"branch={config.branch}, "
f"namespace={config.namespace},"
f"timeout={config.timeout}"
)
try:
# Establish connection to Bauplan
client = bauplan.Client(
api_key=config.api_key,
branch=config.branch,
namespace=config.namespace,
client_timeout=config.timeout
)
logger.info(f"Connected to Bauplan. branch={client.profile.branch} - namespace={client.profile.namespace}")
return client
except Exception as e:
# Handle unexpected errors
logger.error(f"Failed to connect to Bauplan: {str(e)}", exc_info=True)
raise ConnectionError(f"Unable to connect to Bauplan: {str(e)}")
def execute_query(query: str):
# Initialize Bauplan client
bauplan_client = create_bauplan_client()
try:
# Create a response structure optimized for LLM consumption
response = {
"status": "success",
"data": [],
"metadata": {},
"error": None
}
print(f"Executing query: {query}")
# Execute query and get results as Arrow table
result = bauplan_client.query(
query=query,
ref=config.branch,
namespace=config.namespace,
client_timeout=config.timeout
)
# Convert pyarrow.Table to list of dictionaries with native Python values
data_rows = [
dict(zip(result.column_names, [val.as_py() for val in row]))
for row in zip(*[result[col] for col in result.column_names])
]
# Add data and metadata to response
response["data"] = data_rows
response["metadata"] = {
"row_count": len(data_rows),
"column_names": result.column_names,
"column_types": [str(field.type) for field in result.schema],
"query_time": datetime.datetime.now().isoformat(),
"query": query,
}
logger.info(f"Query returned {len(data_rows)} rows")
except Exception as err:
# Consistent error handling with detailed information
error_message = str(err)
logger.error(f"Error executing query: {error_message}")
# Update response for error case
response["status"] = "error"
response["error"] = error_message
response["data"] = [] # Ensure empty data on error
return response
@mcp.tool(name="list_tables", description="List all data tables in the Bauplan data store.")
def list_tables():
"""
List all data tables n the configured Bauplan branch and namespace.
Returns:
dict: Tables object with names
"""
logger.info(f"Executing list_tables")
try:
client = create_bauplan_client()
ret = client.get_tables(ref=config.branch, namespace=config.namespace)
tables = {"tables": [{"name": table.name} for table in ret]}
return tables
except Exception as err:
# Consistent error handling with detailed information
error_message = str(err)
logger.error(f"Error executing get_schema: {error_message}")
response = {
"status": "error",
"error": error_message
}
return response
@mcp.tool(name="get_schema", description="Get the schema of the Bauplan data store.")
def get_schema():
"""
Return the schema of all data tables in the configured Bauplan branch and namespace.
Returns:
dict: Schema object with table fields
"""
logger.info(f"Executing get_schema")
try:
client = create_bauplan_client()
# Get the tables list
ret = client.get_tables(ref=config.branch, namespace=config.namespace)
tables = {"tables": [{"name": table.name} for table in ret]}
# Iterate to get schemas and build final structure
schema_data = {
"schema": [
{
"table": {
"name": t["name"],
"fields": client.get_table(
table=f"{config.namespace}.{t['name']}",
ref=config.branch,
include_raw=True
).raw['schemas'][0]['fields']
}
}
for t in tables["tables"]
]
}
return schema_data
except Exception as err:
# Consistent error handling with detailed information
error_message = str(err)
logger.error(f"Error executing get_schema: {error_message}")
response = {
"status": "error",
"error": error_message
}
return response
@mcp.tool(name="run_query", description="Query the Bauplan data store using SQL.")
def run_query(query: str):
"""
Executes a query against the Bauplan data store with timeout protection.
Args:
query (str): SQL query to execute
Returns:
dict: Response object with status, data, and error information
"""
# Log query for debugging and audit purposes
logger.info(f"Executing query: {query}")
# Enforce SELECT query for security (prevent other operations)
if not query.strip().upper().startswith("SELECT"):
logger.info(f"Exiting: only SELECT queries are permitted")
return {
"status": "error",
"error": "Only SELECT queries are permitted",
"data": [],
"metadata": {"query": query}
}
result = execute_query(query)
return result
