1+ package mcp
2+
3+ import (
4+ "context"
5+ "fmt"
6+
7+ _ "github.com/go-sql-driver/mysql"
8+ "github.com/mark3labs/mcp-go/mcp"
9+ "github.com/mark3labs/mcp-go/server"
10+ "github.com/planetscale/cli/internal/cmdutil"
11+ "github.com/spf13/cobra"
12+ )
13+
14+ // Tool handler function type
15+ type ToolHandler func (ctx context.Context , request mcp.CallToolRequest , ch * cmdutil.Helper ) (* mcp.CallToolResult , error )
16+
17+ // Tool definition
18+ type ToolDef struct {
19+ tool mcp.Tool
20+ handler ToolHandler
21+ }
22+
23+ // getToolDefinitions returns the list of all available MCP tools
24+ func getToolDefinitions () []ToolDef {
25+ return []ToolDef {
26+ {
27+ tool : mcp .NewTool ("list_orgs" ,
28+ mcp .WithDescription ("List all available organizations" ),
29+ ),
30+ handler : HandleListOrgs ,
31+ },
32+ {
33+ tool : mcp .NewTool ("list_databases" ,
34+ mcp .WithDescription ("List all databases in an organization" ),
35+ mcp .WithString ("org" ,
36+ mcp .Description ("The organization name (uses default organization if not specified)" ),
37+ ),
38+ ),
39+ handler : HandleListDatabases ,
40+ },
41+ {
42+ tool : mcp .NewTool ("list_branches" ,
43+ mcp .WithDescription ("List all branches for a database" ),
44+ mcp .WithString ("database" ,
45+ mcp .Description ("The database name" ),
46+ mcp .Required (),
47+ ),
48+ mcp .WithString ("org" ,
49+ mcp .Description ("The organization name (uses default organization if not specified)" ),
50+ ),
51+ ),
52+ handler : HandleListBranches ,
53+ },
54+ {
55+ tool : mcp .NewTool ("list_keyspaces" ,
56+ mcp .WithDescription ("List all keyspaces within a branch" ),
57+ mcp .WithString ("database" ,
58+ mcp .Description ("The database name" ),
59+ mcp .Required (),
60+ ),
61+ mcp .WithString ("branch" ,
62+ mcp .Description ("The branch name" ),
63+ mcp .Required (),
64+ ),
65+ mcp .WithString ("org" ,
66+ mcp .Description ("The organization name (uses default organization if not specified)" ),
67+ ),
68+ ),
69+ handler : HandleListKeyspaces ,
70+ },
71+ {
72+ tool : mcp .NewTool ("list_tables" ,
73+ mcp .WithDescription ("List all tables in a keyspace" ),
74+ mcp .WithString ("database" ,
75+ mcp .Description ("The database name" ),
76+ mcp .Required (),
77+ ),
78+ mcp .WithString ("branch" ,
79+ mcp .Description ("The branch name" ),
80+ mcp .Required (),
81+ ),
82+ mcp .WithString ("keyspace" ,
83+ mcp .Description ("The keyspace name" ),
84+ mcp .Required (),
85+ ),
86+ mcp .WithString ("org" ,
87+ mcp .Description ("The organization name (uses default organization if not specified)" ),
88+ ),
89+ ),
90+ handler : HandleListTables ,
91+ },
92+ {
93+ tool : mcp .NewTool ("get_schema" ,
94+ mcp .WithDescription ("Get the SQL schema for tables in a keyspace" ),
95+ mcp .WithString ("database" ,
96+ mcp .Description ("The database name" ),
97+ mcp .Required (),
98+ ),
99+ mcp .WithString ("branch" ,
100+ mcp .Description ("The branch name" ),
101+ mcp .Required (),
102+ ),
103+ mcp .WithString ("keyspace" ,
104+ mcp .Description ("The keyspace name" ),
105+ mcp .Required (),
106+ ),
107+ mcp .WithString ("tables" ,
108+ mcp .Description ("Tables to get schemas for (single name, comma-separated list, or '*' for all tables)" ),
109+ mcp .Required (),
110+ ),
111+ mcp .WithString ("org" ,
112+ mcp .Description ("The organization name (uses default organization if not specified)" ),
113+ ),
114+ ),
115+ handler : HandleGetSchema ,
116+ },
117+ {
118+ tool : mcp .NewTool ("run_query" ,
119+ mcp .WithDescription ("Run a SQL query against a database branch keyspace" ),
120+ mcp .WithString ("database" ,
121+ mcp .Description ("The database name" ),
122+ mcp .Required (),
123+ ),
124+ mcp .WithString ("branch" ,
125+ mcp .Description ("The branch name" ),
126+ mcp .Required (),
127+ ),
128+ mcp .WithString ("keyspace" ,
129+ mcp .Description ("The keyspace name" ),
130+ mcp .Required (),
131+ ),
132+ mcp .WithString ("query" ,
133+ mcp .Description ("The SQL query to run (read-only queries only)" ),
134+ mcp .Required (),
135+ ),
136+ mcp .WithString ("org" ,
137+ mcp .Description ("The organization name (uses default organization if not specified)" ),
138+ ),
139+ ),
140+ handler : HandleRunQuery ,
141+ },
142+ }
143+ }
144+
145+ // ServerCmd returns a new cobra.Command for the mcp server command.
146+ func ServerCmd (ch * cmdutil.Helper ) * cobra.Command {
147+ cmd := & cobra.Command {
148+ Use : "server" ,
149+ Short : "Start the MCP server" ,
150+ Long : `Start the PlanetScale model context protocol (MCP) server.` ,
151+ RunE : func (cmd * cobra.Command , args []string ) error {
152+ // Create a new MCP server
153+ s := server .NewMCPServer (
154+ "PlanetScale MCP Server" ,
155+ "0.1.0" ,
156+ )
157+
158+ // Register all tools
159+ for _ , toolDef := range getToolDefinitions () {
160+ // Create a tool-specific handler that will forward to our function
161+ // We need to create a local copy of the tool definition to avoid closure issues
162+ def := toolDef
163+ // AddTool expects the tool value directly
164+ s .AddTool (def .tool , func (ctx context.Context , request mcp.CallToolRequest ) (* mcp.CallToolResult , error ) {
165+ return def .handler (ctx , request , ch )
166+ })
167+ }
168+
169+ // Start the server
170+ if err := server .ServeStdio (s ); err != nil {
171+ return fmt .Errorf ("MCP server error: %v" , err )
172+ }
173+
174+ return nil
175+ },
176+ }
177+
178+ return cmd
179+ }
0 commit comments