363 lines
11 KiB
Go
363 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"os/signal"
|
|
"path/filepath"
|
|
"syscall"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/registry/internal/model"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
var (
|
|
mongoURI string
|
|
databaseName string
|
|
collectionName string
|
|
seedFilePath string
|
|
dropCollection bool
|
|
)
|
|
|
|
func init() {
|
|
flag.StringVar(&mongoURI, "uri", "mongodb://localhost:27017", "MongoDB connection URI")
|
|
flag.StringVar(&databaseName, "db", "mcp_registry", "MongoDB database name")
|
|
flag.StringVar(&collectionName, "collection", "servers", "MongoDB collection name")
|
|
flag.StringVar(&seedFilePath, "seed", "", "Path to seed.json file (default: data/seed.json)")
|
|
flag.BoolVar(&dropCollection, "drop", false, "Drop collection before importing")
|
|
flag.Parse()
|
|
|
|
// Set default seed file path if not provided
|
|
if seedFilePath == "" {
|
|
// Try to find the seed.json in the data directory
|
|
seedFilePath = filepath.Join("data", "seed.json")
|
|
if _, err := os.Stat(seedFilePath); os.IsNotExist(err) {
|
|
// Try from the repository root
|
|
repoRoot := filepath.Join("..", "..")
|
|
seedFilePath = filepath.Join(repoRoot, "data", "seed.json")
|
|
}
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
// Create a context with timeout for the database operations
|
|
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
|
defer cancel()
|
|
|
|
// Handle graceful shutdown
|
|
sigChan := make(chan os.Signal, 1)
|
|
signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
|
|
go func() {
|
|
<-sigChan
|
|
log.Println("Received termination signal. Cleaning up...")
|
|
cancel()
|
|
os.Exit(0)
|
|
}()
|
|
|
|
// Read the seed file
|
|
seedData, err := readSeedFile(seedFilePath)
|
|
if err != nil {
|
|
log.Fatalf("Failed to read seed file: %v", err)
|
|
}
|
|
|
|
// Connect to MongoDB
|
|
log.Printf("Connecting to MongoDB at %s", mongoURI)
|
|
client, err := mongo.Connect(ctx, options.Client().ApplyURI(mongoURI))
|
|
if err != nil {
|
|
log.Fatalf("Failed to connect to MongoDB: %v", err)
|
|
}
|
|
defer func() {
|
|
if err = client.Disconnect(ctx); err != nil {
|
|
log.Fatalf("Failed to disconnect from MongoDB: %v", err)
|
|
}
|
|
}()
|
|
|
|
// Ping MongoDB to verify the connection
|
|
if err = client.Ping(ctx, nil); err != nil {
|
|
log.Fatalf("Failed to ping MongoDB: %v", err)
|
|
}
|
|
log.Println("Successfully connected to MongoDB")
|
|
|
|
// Get collection
|
|
db := client.Database(databaseName)
|
|
collection := db.Collection(collectionName)
|
|
|
|
// Drop collection if requested
|
|
if dropCollection {
|
|
log.Printf("Dropping collection %s", collectionName)
|
|
if err := collection.Drop(ctx); err != nil {
|
|
log.Fatalf("Failed to drop collection: %v", err)
|
|
}
|
|
}
|
|
|
|
// Import data
|
|
importData(ctx, collection, seedData)
|
|
}
|
|
|
|
// readSeedFile reads and parses the seed.json file
|
|
func readSeedFile(path string) ([]model.ServerDetail, error) {
|
|
log.Printf("Reading seed file from %s", path)
|
|
|
|
// Read the file content
|
|
fileContent, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to read file: %v", err)
|
|
}
|
|
|
|
// Parse the JSON content
|
|
var servers []model.ServerDetail
|
|
if err := json.Unmarshal(fileContent, &servers); err != nil {
|
|
// Try parsing as a raw JSON array and then convert to our model
|
|
var rawData []map[string]interface{}
|
|
if jsonErr := json.Unmarshal(fileContent, &rawData); jsonErr != nil {
|
|
return nil, fmt.Errorf("failed to parse JSON: %v (original error: %v)", jsonErr, err)
|
|
}
|
|
|
|
// Convert raw data to model format
|
|
servers = convertToServerDetails(rawData)
|
|
}
|
|
|
|
log.Printf("Found %d server entries in seed file", len(servers))
|
|
return servers, nil
|
|
}
|
|
|
|
// convertToServerDetails converts raw JSON data to ServerDetail structs
|
|
func convertToServerDetails(rawData []map[string]interface{}) []model.ServerDetail {
|
|
servers := make([]model.ServerDetail, 0, len(rawData))
|
|
|
|
for _, data := range rawData {
|
|
server := model.ServerDetail{}
|
|
|
|
// Extract basic fields
|
|
if id, ok := data["id"].(string); ok {
|
|
server.ID = id
|
|
}
|
|
if name, ok := data["name"].(string); ok {
|
|
server.Name = name
|
|
}
|
|
if desc, ok := data["description"].(string); ok {
|
|
server.Description = desc
|
|
}
|
|
if version, ok := data["version"].(string); ok {
|
|
server.Version = version
|
|
}
|
|
if canonical, ok := data["registrycanonical"].(string); ok {
|
|
server.RegistryCanonical = canonical
|
|
}
|
|
|
|
// Extract repository
|
|
if repo, ok := data["repository"].(map[string]interface{}); ok {
|
|
server.Repository = model.Repository{
|
|
URL: getStringValue(repo, "url"),
|
|
SubFolder: getStringValue(repo, "subfolder"),
|
|
Branch: getStringValue(repo, "branch"),
|
|
Commit: getStringValue(repo, "commit"),
|
|
}
|
|
}
|
|
|
|
// Extract registries
|
|
if registries, ok := data["registries"].([]interface{}); ok {
|
|
for _, reg := range registries {
|
|
if regMap, ok := reg.(map[string]interface{}); ok {
|
|
registry := model.Registries{
|
|
Name: getStringValue(regMap, "name"),
|
|
PackageName: getStringValue(regMap, "packagename"),
|
|
License: getStringValue(regMap, "license"),
|
|
}
|
|
|
|
// Handle command arguments if present
|
|
if cmd, ok := regMap["command_arguments"].(map[string]interface{}); ok {
|
|
commandArgs := model.Command{}
|
|
|
|
// Extract sub commands
|
|
if subCmds, ok := cmd["sub_commands"].([]interface{}); ok {
|
|
for _, sc := range subCmds {
|
|
if scMap, ok := sc.(map[string]interface{}); ok {
|
|
subCmd := model.SubCommand{
|
|
Name: getStringValue(scMap, "name"),
|
|
Description: getStringValue(scMap, "description"),
|
|
}
|
|
|
|
// Extract named arguments in sub commands
|
|
if namedArgs, ok := scMap["named_arguments"].([]interface{}); ok {
|
|
for _, na := range namedArgs {
|
|
if naMap, ok := na.(map[string]interface{}); ok {
|
|
namedArg := extractNamedArgument(naMap)
|
|
subCmd.NamedArguments = append(subCmd.NamedArguments, namedArg)
|
|
}
|
|
}
|
|
}
|
|
|
|
commandArgs.SubCommands = append(commandArgs.SubCommands, subCmd)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Extract positional arguments
|
|
if posArgs, ok := cmd["positional_arguments"].([]interface{}); ok {
|
|
for _, pa := range posArgs {
|
|
if paMap, ok := pa.(map[string]interface{}); ok {
|
|
posArg := model.PositionalArgument{}
|
|
|
|
if pos, ok := paMap["position"].(float64); ok {
|
|
posArg.Position = int(pos)
|
|
}
|
|
|
|
if arg, ok := paMap["argument"].(map[string]interface{}); ok {
|
|
posArg.Argument = extractArgument(arg)
|
|
}
|
|
|
|
commandArgs.PositionalArguments = append(commandArgs.PositionalArguments, posArg)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Extract environment variables
|
|
if envVars, ok := cmd["environment_variables"].([]interface{}); ok {
|
|
for _, ev := range envVars {
|
|
if evMap, ok := ev.(map[string]interface{}); ok {
|
|
envVar := model.EnvironmentVariable{
|
|
Name: getStringValue(evMap, "name"),
|
|
Description: getStringValue(evMap, "description"),
|
|
}
|
|
|
|
if required, ok := evMap["required"].(bool); ok {
|
|
envVar.Required = required
|
|
}
|
|
|
|
commandArgs.EnvironmentVariables = append(commandArgs.EnvironmentVariables, envVar)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Extract named arguments
|
|
if namedArgs, ok := cmd["named_arguments"].([]interface{}); ok {
|
|
for _, na := range namedArgs {
|
|
if naMap, ok := na.(map[string]interface{}); ok {
|
|
namedArg := extractNamedArgument(naMap)
|
|
commandArgs.NamedArguments = append(commandArgs.NamedArguments, namedArg)
|
|
}
|
|
}
|
|
}
|
|
|
|
registry.CommandArguments = commandArgs
|
|
}
|
|
|
|
server.Registries = append(server.Registries, registry)
|
|
}
|
|
}
|
|
}
|
|
|
|
// Extract remotes
|
|
if remotes, ok := data["remotes"].([]interface{}); ok {
|
|
for _, rem := range remotes {
|
|
if remMap, ok := rem.(map[string]interface{}); ok {
|
|
remote := model.Remotes{
|
|
TransportType: getStringValue(remMap, "transport_type"),
|
|
Url: getStringValue(remMap, "url"),
|
|
}
|
|
server.Remotes = append(server.Remotes, remote)
|
|
}
|
|
}
|
|
}
|
|
|
|
servers = append(servers, server)
|
|
}
|
|
|
|
return servers
|
|
}
|
|
|
|
// extractArgument extracts an Argument struct from a map
|
|
func extractArgument(data map[string]interface{}) model.Argument {
|
|
arg := model.Argument{
|
|
Name: getStringValue(data, "name"),
|
|
Description: getStringValue(data, "description"),
|
|
DefaultValue: getStringValue(data, "default_value"),
|
|
}
|
|
|
|
// Extract boolean fields
|
|
if isRequired, ok := data["is_required"].(bool); ok {
|
|
arg.IsRequired = isRequired
|
|
}
|
|
if isEditable, ok := data["is_editable"].(bool); ok {
|
|
arg.IsEditable = isEditable
|
|
}
|
|
if isRepeatable, ok := data["is_repeatable"].(bool); ok {
|
|
arg.IsRepeatable = isRepeatable
|
|
}
|
|
|
|
// Extract string array for choices
|
|
if choices, ok := data["choices"].([]interface{}); ok {
|
|
for _, choice := range choices {
|
|
if strChoice, ok := choice.(string); ok {
|
|
arg.Choices = append(arg.Choices, strChoice)
|
|
}
|
|
}
|
|
}
|
|
|
|
return arg
|
|
}
|
|
|
|
// extractNamedArgument extracts a NamedArguments struct from a map
|
|
func extractNamedArgument(data map[string]interface{}) model.NamedArguments {
|
|
namedArg := model.NamedArguments{
|
|
ShortFlag: getStringValue(data, "short_flag"),
|
|
LongFlag: getStringValue(data, "long_flag"),
|
|
}
|
|
|
|
if requiresValue, ok := data["requires_value"].(bool); ok {
|
|
namedArg.RequiresValue = requiresValue
|
|
}
|
|
|
|
if arg, ok := data["argument"].(map[string]interface{}); ok {
|
|
namedArg.Argument = extractArgument(arg)
|
|
}
|
|
|
|
return namedArg
|
|
}
|
|
|
|
// getStringValue safely extracts a string value from a map
|
|
func getStringValue(data map[string]interface{}, key string) string {
|
|
if val, ok := data[key].(string); ok {
|
|
return val
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// importData imports the seed data into MongoDB
|
|
func importData(ctx context.Context, collection *mongo.Collection, servers []model.ServerDetail) {
|
|
log.Printf("Importing %d servers into collection %s", len(servers), collection.Name())
|
|
|
|
for i, server := range servers {
|
|
// Create filter based on server ID
|
|
filter := bson.M{"id": server.ID}
|
|
|
|
// Create update document
|
|
update := bson.M{"$set": server}
|
|
|
|
// Use upsert to create if not exists or update if exists
|
|
opts := options.Update().SetUpsert(true)
|
|
result, err := collection.UpdateOne(ctx, filter, update, opts)
|
|
if err != nil {
|
|
log.Printf("Error importing server %s: %v", server.ID, err)
|
|
continue
|
|
}
|
|
|
|
if result.UpsertedCount > 0 {
|
|
log.Printf("[%d/%d] Created server: %s", i+1, len(servers), server.Name)
|
|
} else if result.ModifiedCount > 0 {
|
|
log.Printf("[%d/%d] Updated server: %s", i+1, len(servers), server.Name)
|
|
} else {
|
|
log.Printf("[%d/%d] Server already up to date: %s", i+1, len(servers), server.Name)
|
|
}
|
|
}
|
|
|
|
log.Println("Import completed successfully")
|
|
}
|