forge-tools-ssh/internal/tools/db.go

package tools

import (
	"context"
	"fmt"

	"github.com/modelcontextprotocol/go-sdk/mcp"
)

func registerDBTools(s *mcp.Server) {
	// db_query
	s.AddTool(&mcp.Tool{
		Name:        "db_query",
		Description: "Execute SQL/CQL/MongoDB query inside a database container",
		InputSchema: schema(map[string]any{
			"container": strProps("Docker container name running the database", true),
			"db_type":   strProps("Database engine type", true, "postgres", "mysql", "scylladb", "cassandra", "mongodb"),
			"query":     strProps("Query to execute", true),
			"database":  strProps("Database/keyspace name", false),
			"username":  strProps("Database username", false),
			"password":  strProps("Database password", false),
			"timeout":   intProps("Query timeout in seconds (default: 60)", false),
			"target":    strProps("Connection alias (default: primary)", false),
		}, []string{"container", "db_type", "query"}),
	}, dbQueryHandler)

	// db_schema
	s.AddTool(&mcp.Tool{
		Name:        "db_schema",
		Description: "Get database schema (tables/collections list)",
		InputSchema: schema(map[string]any{
			"container": strProps("Docker container name", true),
			"db_type":   strProps("Database engine type", true, "postgres", "mysql", "scylladb", "cassandra", "mongodb"),
			"database":  strProps("Database/keyspace name", false),
			"username":  strProps("Database username", false),
			"password":  strProps("Database password", false),
			"target":    strProps("Connection alias (default: primary)", false),
		}, []string{"container", "db_type"}),
	}, dbSchemaHandler)

	// list_db_containers
	s.AddTool(&mcp.Tool{
		Name:        "list_db_containers",
		Description: "Find Docker containers that look like databases",
		InputSchema: schema(map[string]any{
			"target": strProps("Connection alias (default: primary)", false),
		}, nil),
	}, listDBContainersHandler)
}

func dbQueryHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	m := manager()
	if m == nil {
		return errorResult("no SSH manager initialized"), nil
	}
	args := requestArgs(req)
	container, _ := requireString(args, "container")
	dbType, _ := requireString(args, "db_type")
	query, _ := requireString(args, "query")
	database := getString(args, "database", "")
	username := getString(args, "username", "")
	password := getString(args, "password", "")
	timeout := getInt(args, "timeout", 60)
	target := getString(args, "target", "primary")

	if err := checkDockerAvailable(ctx, m, target); err != nil {
		return errorResult(err.Error()), nil
	}

	var cmd string
	switch dbType {
	case "postgres":
		user := "postgres"
		if username != "" {
			user = username
		}
		db := database
		if db == "" {
			db = "postgres"
		}
		cmd = fmt.Sprintf("docker exec -e PGPASSWORD=%s %s timeout %d psql -U %s -d %s -c %s 2>&1",
			shellQuote(password), shellQuote(container), timeout, shellQuote(user), shellQuote(db), shellQuote(query))

	case "mysql":
		user := "root"
		if username != "" {
			user = username
		}
		cmd = fmt.Sprintf("docker exec %s timeout %d mysql -u%s", shellQuote(container), timeout, shellQuote(user))
		if password != "" {
			cmd += fmt.Sprintf(" -p%s", shellQuote(password))
		}
		if database != "" {
			cmd += fmt.Sprintf(" %s", shellQuote(database))
		}
		cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))

	case "scylladb", "cassandra":
		cmd = fmt.Sprintf("docker exec %s timeout %d cqlsh", shellQuote(container), timeout)
		if username != "" {
			cmd += fmt.Sprintf(" -u %s", shellQuote(username))
		}
		if password != "" {
			cmd += fmt.Sprintf(" -p %s", shellQuote(password))
		}
		cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))

	case "mongodb":
		db := database
		if db == "" {
			db = "admin"
		}
		cmd = fmt.Sprintf("docker exec %s timeout %d mongosh --quiet %s", shellQuote(container), timeout, shellQuote(db))
		if username != "" && password != "" {
			cmd += fmt.Sprintf(" -u %s -p %s --authenticationDatabase admin", shellQuote(username), shellQuote(password))
		}
		cmd += fmt.Sprintf(" --eval %s 2>&1", shellQuote(query))

	default:
		return errorResult(fmt.Sprintf("Unsupported database type: %s. Supported: postgres, mysql, scylladb, cassandra, mongodb", dbType)), nil
	}

	output, err := m.Execute(ctx, cmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(output), nil
}

func dbSchemaHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	m := manager()
	if m == nil {
		return errorResult("no SSH manager initialized"), nil
	}
	args := requestArgs(req)
	container, _ := requireString(args, "container")
	dbType, _ := requireString(args, "db_type")
	database := getString(args, "database", "")
	username := getString(args, "username", "")
	password := getString(args, "password", "")
	target := getString(args, "target", "primary")

	if err := checkDockerAvailable(ctx, m, target); err != nil {
		return errorResult(err.Error()), nil
	}

	var query string
	switch dbType {
	case "postgres":
		query = "\\dt"
	case "mysql":
		query = "SHOW TABLES;"
	case "scylladb", "cassandra":
		if database != "" {
			query = fmt.Sprintf("DESCRIBE KEYSPACE %s;", database)
		} else {
			query = "DESCRIBE KEYSPACES;"
		}
	case "mongodb":
		query = "db.getCollectionNames()"
	default:
		return errorResult(fmt.Sprintf("Unsupported database type: %s", dbType)), nil
	}

	var cmd string
	switch dbType {
	case "postgres":
		user := "postgres"
		if username != "" {
			user = username
		}
		db := database
		if db == "" {
			db = "postgres"
		}
		cmd = fmt.Sprintf("docker exec -e PGPASSWORD=%s %s psql -U %s -d %s -c %s 2>&1",
			shellQuote(password), shellQuote(container), shellQuote(user), shellQuote(db), shellQuote(query))
	case "mysql":
		user := "root"
		if username != "" {
			user = username
		}
		cmd = fmt.Sprintf("docker exec %s mysql -u%s", shellQuote(container), shellQuote(user))
		if password != "" {
			cmd += fmt.Sprintf(" -p%s", shellQuote(password))
		}
		if database != "" {
			cmd += fmt.Sprintf(" %s", shellQuote(database))
		}
		cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))
	case "scylladb", "cassandra":
		cmd = fmt.Sprintf("docker exec %s cqlsh", shellQuote(container))
		if username != "" {
			cmd += fmt.Sprintf(" -u %s", shellQuote(username))
		}
		if password != "" {
			cmd += fmt.Sprintf(" -p %s", shellQuote(password))
		}
		cmd += fmt.Sprintf(" -e %s 2>&1", shellQuote(query))
	case "mongodb":
		db := database
		if db == "" {
			db = "admin"
		}
		cmd = fmt.Sprintf("docker exec %s mongosh --quiet %s", shellQuote(container), shellQuote(db))
		if username != "" && password != "" {
			cmd += fmt.Sprintf(" -u %s -p %s --authenticationDatabase admin", shellQuote(username), shellQuote(password))
		}
		cmd += fmt.Sprintf(" --eval %s 2>&1", shellQuote(query))
	}

	output, err := m.Execute(ctx, cmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	return textResult(output), nil
}

func listDBContainersHandler(ctx context.Context, req *mcp.CallToolRequest) (*mcp.CallToolResult, error) {
	m := manager()
	if m == nil {
		return errorResult("no SSH manager initialized"), nil
	}
	target := getString(requestArgs(req), "target", "primary")

	if err := checkDockerAvailable(ctx, m, target); err != nil {
		return errorResult(err.Error()), nil
	}

	cmd := `docker ps --format '{{.Names}}|{{.Image}}' | while read line; do
  name=$(echo "$line" | cut -d'|' -f1)
  image=$(echo "$line" | cut -d'|' -f2)
  case "$image" in
    *postgres*) echo "$name|$image|postgres" ;;
    *mysql*|*mariadb*) echo "$name|$image|mysql" ;;
    *scylla*) echo "$name|$image|scylladb" ;;
    *cassandra*) echo "$name|$image|cassandra" ;;
    *mongo*) echo "$name|$image|mongodb" ;;
    *redis*) echo "$name|$image|redis" ;;
  esac
done 2>/dev/null`

	output, err := m.Execute(ctx, cmd, target)
	if err != nil {
		return errorResult(err.Error()), nil
	}
	if trimOutput(output) == "" {
		return textResult("No database containers found"), nil
	}
	return textResult(output), nil
}