Add warehouse_id parameter to SQL execution for improved flexibility

Author: characat0

main.go

@@ -60,6 +60,7 @@ func main() {
 		mcp.WithString("statement", mcp.Description("SQL statement to execute"), mcp.Required()),
 		mcp.WithNumber("execution_timeout_seconds", mcp.Description("Maximum time in seconds to wait for query execution"), mcp.DefaultNumber(60)),
 		mcp.WithNumber("max_rows", mcp.Description("Maximum number of rows to return in the result"), mcp.DefaultNumber(100)),
+		mcp.WithString("warehouse_id", mcp.Description("ID of the warehouse to use for execution. If not specified, the first available warehouse will be used")),
 	), tools.WithWorkspaceClientHandler(w, tools.ExecuteTool(tools.ExecuteSQL)))
 
 	s.AddTool(mcp.NewTool("list_warehouses",

tools/execute_sql.go

@@ -20,6 +20,7 @@ func ExecuteSQL(ctx context.Context, request mcp.CallToolRequest) (interface{},
 	sqlStatement := ExtractStringParam(request, "statement", "")
 	timeoutSeconds := ExtractFloatParam(request, "execution_timeout_seconds", 60)
 	maxRows := ExtractFloatParam(request, "max_rows", 100)
+	warehouseId := ExtractStringParam(request, "warehouse_id", "")
 
 	// Poll every 10 seconds
 	if timeoutSeconds < 5 {
@@ -30,21 +31,25 @@ func ExecuteSQL(ctx context.Context, request mcp.CallToolRequest) (interface{},
 	// Poll for statement completion
 	maxAttempts := int(timeoutSeconds / 10)
 
-	// Get available warehouses
-	warehouses, err := w.Warehouses.ListAll(ctx, sql.ListWarehousesRequest{})
-	if err != nil {
-		return nil, fmt.Errorf("error listing warehouses: %w", err)
-	}
-	if len(warehouses) == 0 {
-		return nil, fmt.Errorf("no warehouses available")
+	// Determine which warehouse to use
+	if warehouseId == "" {
+		// Get available warehouses and use the first one
+		warehouses, err := w.Warehouses.ListAll(ctx, sql.ListWarehousesRequest{})
+		if err != nil {
+			return nil, fmt.Errorf("error listing warehouses: %w", err)
+		}
+		if len(warehouses) == 0 {
+			return nil, fmt.Errorf("no warehouses available")
+		}
+		warehouseId = warehouses[0].Id
 	}
 
 	// Execute the SQL statement with the specified row limit
 	res, err := w.StatementExecution.ExecuteStatement(ctx, sql.ExecuteStatementRequest{
 		RowLimit:    int64(maxRows),
 		Statement:   sqlStatement,
 		WaitTimeout: "5s",
-		WarehouseId: warehouses[0].Id,
+		WarehouseId: warehouseId,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("error executing SQL statement: %w", err)