DataPilot MCP Server

Instructions

Implementation Reference

• src/main.py:73-102 (handler)

  MCP tool handler for execute_sql: executes SQL query via SnowflakeClient, handles errors, provides context logging
  @mcp.tool() async def execute_sql(request: SQLQueryRequest, ctx: Context) -> QueryResult: """Execute a SQL query on Snowflake and return results""" await ctx.info(f"Executing SQL query: {request.query[:100]}...") try: client = await get_snowflake_client() result = await client.execute_query( request.query, request.limit, request.warehouse ) if result.success: await ctx.info(f"Query executed successfully. {result.row_count} rows returned.") else: await ctx.error(f"Query failed: {result.error}") return result except Exception as e: logger.error(f"Error executing SQL: {str(e)}") await ctx.error(f"Failed to execute query: {str(e)}") return QueryResult( success=False, data=[], columns=[], row_count=0, error=str(e) )
• src/models.py:17-22 (schema)

  Input schema for execute_sql tool: defines query, optional limit and warehouse
  class SQLQueryRequest(BaseModel): """Request model for SQL query execution""" query: str = Field(..., description="SQL query to execute") limit: Optional[int] = Field(None, description="Maximum number of rows to return") warehouse: Optional[str] = Field(None, description="Warehouse to use for this query")
• src/models.py:44-53 (schema)

  Output schema for execute_sql tool: query results including data, columns, metadata, and error handling
  class QueryResult(BaseModel): """Result of a SQL query execution""" success: bool data: List[Dict[str, Any]] columns: List[str] row_count: int execution_time_ms: Optional[int] = None query_id: Optional[str] = None warehouse_used: Optional[str] = None error: Optional[str] = None
• src/main.py:73-73 (registration)

  MCP tool registration decorator for execute_sql
  @mcp.tool()
• src/snowflake_client.py:58-118 (helper)

  Supporting helper: SnowflakeClient.execute_query method that performs the actual database query execution and formats results
  async def execute_query( self, query: str, limit: Optional[int] = None, warehouse: Optional[str] = None ) -> QueryResult: """Execute a SQL query and return results""" start_time = time.time() try: # Switch warehouse if specified if warehouse: await self._use_warehouse(warehouse) # Add limit if specified if limit and not query.strip().upper().startswith('SELECT'): logger.warning("LIMIT can only be applied to SELECT queries") elif limit and 'LIMIT' not in query.upper(): query = f"{query.rstrip(';')} LIMIT {limit}" async with self.get_cursor() as cursor: logger.info(f"Executing query: {query[:100]}...") await asyncio.to_thread(cursor.execute, query) results = await asyncio.to_thread(cursor.fetchall) # Get column names columns = [desc[0] for desc in cursor.description] if cursor.description else [] # Convert results to list of dictionaries data = [] for row in results: if isinstance(row, dict): data.append(row) else: data.append(dict(zip(columns, row))) execution_time_ms = int((time.time() - start_time) * 1000) return QueryResult( success=True, data=data, columns=columns, row_count=len(data), execution_time_ms=execution_time_ms, query_id=cursor.sfqid if hasattr(cursor, 'sfqid') else None, warehouse_used=warehouse or self.config.warehouse ) except Exception as e: logger.error(f"Query execution failed: {str(e)}") execution_time_ms = int((time.time() - start_time) * 1000) return QueryResult( success=False, data=[], columns=[], row_count=0, execution_time_ms=execution_time_ms, error=str(e) )