Merge pull request #1105 from planetscale/piki/mcp-postgres

Implement all the MCP tools for Postgres

Author: piki

go.mod

@@ -80,6 +80,7 @@ require (
 	github.com/klauspost/connect-compress/v2 v2.0.0 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
+	github.com/lib/pq v1.10.9 // indirect
 	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect

go.sum

@@ -128,6 +128,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7 h1:k/1ku0yehLCPqERCHkIHMDqDg1R02AcCScRuHbamU3s=
 github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7/go.mod h1:YR/zYthNdWfO8+0IOyHDcIDBBBS2JMnYUIwSsnwmRqU=
+github.com/lib/pq v1.10.9 h1:YXG7RB+JIjhP29X+OtkiDnYaXQwpS4JEWq7dtCCRUEw=
+github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
 github.com/lucasb-eyer/go-colorful v1.2.0 h1:1nnpGOrhyZZuNyfu1QjKiUICQ74+3FNCN69Aj6K7nkY=
 github.com/lucasb-eyer/go-colorful v1.2.0/go.mod h1:R4dSotOR9KMtayYi1e77YzuveK+i7ruzyGqttikkLy0=
 github.com/mailru/easyjson v0.7.7 h1:UGYAvKxe3sBsEDzO8ZeWOSlIQfWFlxbzLZe7hwFURr0=

internal/cmd/mcp/server.go

@@ -22,6 +22,7 @@ type ToolDef struct {
 
 // getToolDefinitions returns the list of all available MCP tools
 func getToolDefinitions() []ToolDef {
+	namingBlurb := ".  Two common naming conventions for PlanetScale databases are <org>/<database> and <org>/<database>/<branch>. When the user provides a database identifier in either of these formats, automatically parse and use the org, database, and branch parameters directly - do not perform discovery steps like list_orgs or list_databases.  Examples: `acme/widgets` -> org=acme, database=widgets.  `acme/widgets/main` -> org=acme, database=widgets, branch=main.  If the user provides an identifier like 'org/database' or 'org/database/branch', parse these components directly and skip organizational/database discovery steps."
 	return []ToolDef{
 		{
 			tool: mcp.NewTool("list_orgs",
@@ -33,27 +34,29 @@ func getToolDefinitions() []ToolDef {
 			tool: mcp.NewTool("list_databases",
 				mcp.WithDescription("List all databases in an organization"),
 				mcp.WithString("org",
-					mcp.Description("The organization name (uses default organization if not specified)"),
+					mcp.Description("The organization name"),
+					mcp.Required(),
 				),
 			),
 			handler: HandleListDatabases,
 		},
 		{
 			tool: mcp.NewTool("list_branches",
-				mcp.WithDescription("List all branches for a database"),
+				mcp.WithDescription("List all branches for a database"+namingBlurb),
 				mcp.WithString("database",
 					mcp.Description("The database name"),
 					mcp.Required(),
 				),
 				mcp.WithString("org",
-					mcp.Description("The organization name (uses default organization if not specified)"),
+					mcp.Description("The organization name"),
+					mcp.Required(),
 				),
 			),
 			handler: HandleListBranches,
 		},
 		{
 			tool: mcp.NewTool("list_keyspaces",
-				mcp.WithDescription("List all keyspaces within a branch"),
+				mcp.WithDescription("List all keyspaces within a branch"+namingBlurb),
 				mcp.WithString("database",
 					mcp.Description("The database name"),
 					mcp.Required(),
@@ -63,14 +66,15 @@ func getToolDefinitions() []ToolDef {
 					mcp.Required(),
 				),
 				mcp.WithString("org",
-					mcp.Description("The organization name (uses default organization if not specified)"),
+					mcp.Description("The organization name"),
+					mcp.Required(),
 				),
 			),
 			handler: HandleListKeyspaces,
 		},
 		{
 			tool: mcp.NewTool("list_tables",
-				mcp.WithDescription("List all tables in a keyspace"),
+				mcp.WithDescription("List all tables in a keyspace/database"+namingBlurb),
 				mcp.WithString("database",
 					mcp.Description("The database name"),
 					mcp.Required(),
@@ -80,18 +84,21 @@ func getToolDefinitions() []ToolDef {
 					mcp.Required(),
 				),
 				mcp.WithString("keyspace",
-					mcp.Description("The keyspace name"),
-					mcp.Required(),
+					mcp.Description("The keyspace name (for MySQL) or inner database name (for PostgreSQL)"),
+				),
+				mcp.WithString("schema",
+					mcp.Description("Filter tables by schema (PostgreSQL only)"),
 				),
 				mcp.WithString("org",
-					mcp.Description("The organization name (uses default organization if not specified)"),
+					mcp.Description("The organization name"),
+					mcp.Required(),
 				),
 			),
 			handler: HandleListTables,
 		},
 		{
 			tool: mcp.NewTool("get_schema",
-				mcp.WithDescription("Get the SQL schema for tables in a keyspace"),
+				mcp.WithDescription("Get the SQL schema for tables in a keyspace/database"+namingBlurb),
 				mcp.WithString("database",
 					mcp.Description("The database name"),
 					mcp.Required(),
@@ -101,22 +108,22 @@ func getToolDefinitions() []ToolDef {
 					mcp.Required(),
 				),
 				mcp.WithString("keyspace",
-					mcp.Description("The keyspace name"),
-					mcp.Required(),
+					mcp.Description("The keyspace name (for MySQL) or inner database name (for PostgreSQL)"),
 				),
 				mcp.WithString("tables",
-					mcp.Description("Tables to get schemas for (single name, comma-separated list, or '*' for all tables)"),
+					mcp.Description("Tables to get schemas for. MySQL: comma-separated list of table names, or '*' for all tables. PostgreSQL: comma-separated list of simple table names in the public schema, qualified schema.table_names, 'schema.*' for all tables in a schema, or '*' for all tables in all schemas"),
 					mcp.Required(),
 				),
 				mcp.WithString("org",
-					mcp.Description("The organization name (uses default organization if not specified)"),
+					mcp.Description("The organization name"),
+					mcp.Required(),
 				),
 			),
 			handler: HandleGetSchema,
 		},
 		{
 			tool: mcp.NewTool("run_query",
-				mcp.WithDescription("Run a SQL query against a database branch keyspace"),
+				mcp.WithDescription("Run a SQL query against a database branch keyspace/database"+namingBlurb),
 				mcp.WithString("database",
 					mcp.Description("The database name"),
 					mcp.Required(),
@@ -126,22 +133,22 @@ func getToolDefinitions() []ToolDef {
 					mcp.Required(),
 				),
 				mcp.WithString("keyspace",
-					mcp.Description("The keyspace name"),
-					mcp.Required(),
+					mcp.Description("The keyspace name (for MySQL) or inner database name (for PostgreSQL)"),
 				),
 				mcp.WithString("query",
 					mcp.Description("The SQL query to run (read-only queries only)"),
 					mcp.Required(),
 				),
 				mcp.WithString("org",
-					mcp.Description("The organization name (uses default organization if not specified)"),
+					mcp.Description("The organization name"),
+					mcp.Required(),
 				),
 			),
 			handler: HandleRunQuery,
 		},
 		{
 			tool: mcp.NewTool("get_insights",
-				mcp.WithDescription("Get recent performance data for a database branch"),
+				mcp.WithDescription("Get recent performance data for a database branch"+namingBlurb),
 				mcp.WithString("database",
 					mcp.Description("The database name"),
 					mcp.Required(),
@@ -151,7 +158,8 @@ func getToolDefinitions() []ToolDef {
 					mcp.Required(),
 				),
 				mcp.WithString("org",
-					mcp.Description("The organization name (uses default organization if not specified)"),
+					mcp.Description("The organization name"),
+					mcp.Required(),
 				),
 			),
 			handler: HandleGetInsights,
@@ -170,6 +178,20 @@ func ServerCmd(ch *cmdutil.Helper) *cobra.Command {
 			s := server.NewMCPServer(
 				"PlanetScale MCP Server",
 				"0.1.0",
+				server.WithInstructions(`PlanetScale Database Management tools
+
+These tools provide read-only access to MySQL (Vitess) and PostgreSQL
+databases hosted by PlanetScale.  The naming convention for PlanetScale
+databases is <org>/<database> or <org>/<database>/<branch>, so if the user
+refers to databases in this format, they are specifying the org, database,
+and possibly branch needed to invoke the tools in this MCP.  Usually if
+the branch is omitted, the default is "main".
+
+When PlanetScale uses the word "database", it refers to a collection of
+database branches (production, development, staging, etc.), each with one
+or more replicas.  This is distinct from PostgreSQL's notion of a database
+as a namespace on a single server, which PlanetScale calls a "keyspace"
+for both MySQL and PostgreSQL.`),
 			)
 
 			// Register all tools