From 6b22cf09b376ed94c772c69951a2d125e495b26b Mon Sep 17 00:00:00 2001 From: Yusheng Guo <113574546+yshngg@users.noreply.github.com> Date: Wed, 4 Jun 2025 21:26:48 +0800 Subject: [PATCH] feat(config): specify to use in-memory database (#74) Add support for specifying database type (MongoDB or in-memory database). ## Motivation and Context The README states that MCP Registry supports both MongoDB and in-memory databases, but does not provide the capability to specify which database type to use. https://github.com/modelcontextprotocol/registry/blob/5d9171bdccf579398273806a7067cda29e88363c/README.md?plain=1#L19 This PR addresses this limitation by adding explicit database type configuration through the `DATABASE_TYPE` environment variable, allowing users to choose between `mongodb` and `memory` database types in a standardized way. ## How Has This Been Tested? Yes ```console $ export MCP_REGISTRY_DATABASE_TYPE="memory" $ ./scripts/test_endpoints.sh Testing health endpoint: http://localhost:8080/v0/health Status Code: 200 Response: { "status": "ok", "github_client_id": "" } Health check successful ------------------------------------- Testing servers endpoint: http://localhost:8080/v0/servers Status Code: 200 Response Summary: Total registries: 0 servers Names: Pagination Metadata: {} servers Details: { "servers": [], "metadata": {} } servers request successful ------------------------------------- Testing ping endpoint: http://localhost:8080/v0/ping Status Code: 200 Response: { "status": "ok", "version": "dev" } Ping successful ------------------------------------- All tests passed successfully! ``` ## Breaking Changes No ## Types of changes - [ ] Bug fix (non-breaking change which fixes an issue) - [x] New feature (non-breaking change which adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [x] Documentation update ## Checklist - [x] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [x] My code follows the repository's style guidelines - [x] New and existing tests pass locally - [x] I have added appropriate error handling - [x] I have added or updated documentation as needed ## Additional context Close #63 --------- Co-authored-by: Avinash Sridhar Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- README.md | 1 + cmd/registry/main.go | 71 ++++++++++++++++++++++----------------- internal/config/config.go | 28 +++++++++------ 3 files changed, 60 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 5cd446a..9ac3cdb 100644 --- a/README.md +++ b/README.md @@ -289,6 +289,7 @@ The service can be configured using environment variables: | Variable | Description | Default | |----------|-------------|---------| | `MCP_REGISTRY_APP_VERSION` | Application version | `dev` | +| `MCP_REGISTRY_DATABASE_TYPE` | Database type | `mongodb` | | `MCP_REGISTRY_COLLECTION_NAME` | MongoDB collection name | `servers_v2` | | `MCP_REGISTRY_DATABASE_NAME` | MongoDB database name | `mcp-registry` | | `MCP_REGISTRY_DATABASE_URL` | MongoDB connection string | `mongodb://localhost:27017` | diff --git a/cmd/registry/main.go b/cmd/registry/main.go index 7042b51..cb776f5 100644 --- a/cmd/registry/main.go +++ b/cmd/registry/main.go @@ -15,6 +15,7 @@ import ( "github.com/modelcontextprotocol/registry/internal/auth" "github.com/modelcontextprotocol/registry/internal/config" "github.com/modelcontextprotocol/registry/internal/database" + "github.com/modelcontextprotocol/registry/internal/model" "github.com/modelcontextprotocol/registry/internal/service" ) @@ -39,40 +40,50 @@ func main() { // Initialize services based on environment var registryService service.RegistryService - // Use MongoDB for real registry service in production/other environments - // Create a context with timeout for MongoDB connection - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() + switch cfg.DatabaseType { + case config.DatabaseTypeMemory: + memoryDB := database.NewMemoryDB(map[string]*model.Server{}) + registryService = service.NewRegistryServiceWithDB(memoryDB) + case config.DatabaseTypeMongoDB: + // Use MongoDB for real registry service in production/other environments + // Create a context with timeout for MongoDB connection + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() - // Connect to MongoDB - mongoDB, err := database.NewMongoDB(ctx, cfg.DatabaseURL, cfg.DatabaseName, cfg.CollectionName) - if err != nil { - log.Printf("Failed to connect to MongoDB: %v", err) + // Connect to MongoDB + mongoDB, err := database.NewMongoDB(ctx, cfg.DatabaseURL, cfg.DatabaseName, cfg.CollectionName) + if err != nil { + log.Printf("Failed to connect to MongoDB: %v", err) + return + } + + // Create registry service with MongoDB + registryService = service.NewRegistryServiceWithDB(mongoDB) + log.Printf("MongoDB database name: %s", cfg.DatabaseName) + log.Printf("MongoDB collection name: %s", cfg.CollectionName) + + // Store the MongoDB instance for later cleanup + defer func() { + if err := mongoDB.Close(); err != nil { + log.Printf("Error closing MongoDB connection: %v", err) + } else { + log.Println("MongoDB connection closed successfully") + } + }() + + if cfg.SeedImport { + log.Println("Importing data...") + if err := database.ImportSeedFile(mongoDB, cfg.SeedFilePath); err != nil { + log.Printf("Failed to import seed file: %v", err) + } else { + log.Println("Data import completed successfully") + } + } + default: + log.Printf("Invalid database type: %s; supported types: %s, %s", cfg.DatabaseType, config.DatabaseTypeMemory, config.DatabaseTypeMongoDB) return } - // Create registry service with MongoDB - registryService = service.NewRegistryServiceWithDB(mongoDB) - log.Printf("MongoDB database name: %s", cfg.DatabaseName) - log.Printf("MongoDB collection name: %s", cfg.CollectionName) - - // Store the MongoDB instance for later cleanup - defer func() { - if err := mongoDB.Close(); err != nil { - log.Printf("Error closing MongoDB connection: %v", err) - } else { - log.Println("MongoDB connection closed successfully") - } - }() - - if cfg.SeedImport { - log.Println("Importing data...") - if err := database.ImportSeedFile(mongoDB, cfg.SeedFilePath); err != nil { - log.Printf("Failed to import seed file: %v", err) - } - log.Println("Data import completed successfully") - } - // Initialize authentication services authService := auth.NewAuthService(cfg) diff --git a/internal/config/config.go b/internal/config/config.go index cbb5637..950c3f5 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -4,18 +4,26 @@ import ( env "github.com/caarlos0/env/v11" ) +type DatabaseType string + +const ( + DatabaseTypeMongoDB DatabaseType = "mongodb" + DatabaseTypeMemory DatabaseType = "memory" +) + // Config holds the application configuration type Config struct { - ServerAddress string `env:"SERVER_ADDRESS" envDefault:":8080"` - DatabaseURL string `env:"DATABASE_URL" envDefault:"mongodb://localhost:27017"` - DatabaseName string `env:"DATABASE_NAME" envDefault:"mcp-registry"` - CollectionName string `env:"COLLECTION_NAME" envDefault:"servers_v2"` - LogLevel string `env:"LOG_LEVEL" envDefault:"info"` - SeedFilePath string `env:"SEED_FILE_PATH" envDefault:"data/seed.json"` - SeedImport bool `env:"SEED_IMPORT" envDefault:"true"` - Version string `env:"VERSION" envDefault:"dev"` - GithubClientID string `env:"GITHUB_CLIENT_ID" envDefault:""` - GithubClientSecret string `env:"GITHUB_CLIENT_SECRET" envDefault:""` + ServerAddress string `env:"SERVER_ADDRESS" envDefault:":8080"` + DatabaseType DatabaseType `env:"DATABASE_TYPE" envDefault:"mongodb"` + DatabaseURL string `env:"DATABASE_URL" envDefault:"mongodb://localhost:27017"` + DatabaseName string `env:"DATABASE_NAME" envDefault:"mcp-registry"` + CollectionName string `env:"COLLECTION_NAME" envDefault:"servers_v2"` + LogLevel string `env:"LOG_LEVEL" envDefault:"info"` + SeedFilePath string `env:"SEED_FILE_PATH" envDefault:"data/seed.json"` + SeedImport bool `env:"SEED_IMPORT" envDefault:"true"` + Version string `env:"VERSION" envDefault:"dev"` + GithubClientID string `env:"GITHUB_CLIENT_ID" envDefault:""` + GithubClientSecret string `env:"GITHUB_CLIENT_SECRET" envDefault:""` } // NewConfig creates a new configuration with default values