internal/mcpgateway/store_mongodb.go
Rama main · réplica de lectura del repositorio original.
package mcpgateway
import (
"context"
"errors"
"fmt"
"strings"
"time"
"go.mongodb.org/mongo-driver/v2/bson"
"go.mongodb.org/mongo-driver/v2/mongo"
"go.mongodb.org/mongo-driver/v2/mongo/options"
"github.com/enterpilot/gomodel/internal/storage"
)
type mongoMCPServerDocument struct {
ID string `bson:"_id"`
DisplayName string `bson:"display_name,omitempty"`
URL string `bson:"url,omitempty"`
Transport string `bson:"transport,omitempty"`
Headers map[string]string `bson:"headers,omitempty"`
Description string `bson:"description,omitempty"`
Enabled bool `bson:"enabled"`
AllowedTools []string `bson:"allowed_tools,omitempty"`
DisallowedTools []string `bson:"disallowed_tools,omitempty"`
UserPaths []string `bson:"user_paths,omitempty"`
DisallowedUserPaths []string `bson:"disallowed_user_paths,omitempty"`
ToolTimeoutSeconds int `bson:"tool_timeout_seconds,omitempty"`
CreatedAt time.Time `bson:"created_at"`
UpdatedAt time.Time `bson:"updated_at"`
}
type mongoMCPServerIDFilter struct {
ID string `bson:"_id"`
}
// MongoDBStore stores managed MCP servers and virtual servers in MongoDB.
type MongoDBStore struct {
collection *mongo.Collection
virtuals *mongo.Collection
oauth *mongo.Collection
}
// NewMongoDBStore creates collection indexes if needed.
func NewMongoDBStore(database *mongo.Database) (*MongoDBStore, error) {
if database == nil {
return nil, fmt.Errorf("database is required")
}
coll := database.Collection("mcp_servers")
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
indexes := []mongo.IndexModel{
{Keys: bson.D{{Key: "enabled", Value: 1}}},
{Keys: bson.D{{Key: "updated_at", Value: -1}}},
}
if _, err := coll.Indexes().CreateMany(ctx, indexes); err != nil {
return nil, fmt.Errorf("create mcp_servers indexes: %w", err)
}
return &MongoDBStore{
collection: coll,
virtuals: database.Collection("mcp_virtual_servers"),
oauth: database.Collection("mcp_server_oauth"),
}, nil
}
func (s *MongoDBStore) List(ctx context.Context) ([]ManagedServer, error) {
return listMongo(ctx, s.collection, "mcp servers", managedServerFromMongo)
}
// listMongo reads a whole collection sorted by _id, converting each document.
func listMongo[D, T any](ctx context.Context, coll *mongo.Collection, label string, convert func(D) T) ([]T, error) {
cursor, err := coll.Find(ctx, bson.M{}, options.Find().SetSort(bson.D{{Key: "_id", Value: 1}}))
if err != nil {
return nil, fmt.Errorf("list %s: %w", label, err)
}
defer cursor.Close(ctx)
result := make([]T, 0)
for cursor.Next(ctx) {
var doc D
if err := cursor.Decode(&doc); err != nil {
return nil, fmt.Errorf("decode %s: %w", label, err)
}
result = append(result, convert(doc))
}
if err := cursor.Err(); err != nil {
return nil, fmt.Errorf("iterate %s: %w", label, err)
}
return result, nil
}
func (s *MongoDBStore) Get(ctx context.Context, name string) (*ManagedServer, error) {
var doc mongoMCPServerDocument
err := s.collection.FindOne(ctx, mongoMCPServerIDFilter{ID: strings.TrimSpace(name)}).Decode(&doc)
if err != nil {
if errors.Is(err, mongo.ErrNoDocuments) {
return nil, ErrNotFound
}
return nil, fmt.Errorf("get mcp server: %w", err)
}
server := managedServerFromMongo(doc)
return &server, nil
}
func (s *MongoDBStore) Upsert(ctx context.Context, server ManagedServer) error {
stampUpsert(&server)
update := bson.M{
"$set": mongoMCPServerFields(server),
"$setOnInsert": bson.M{
"created_at": server.CreatedAt,
},
}
_, err := s.collection.UpdateOne(ctx, mongoMCPServerIDFilter{ID: strings.TrimSpace(server.Name)}, update, options.UpdateOne().SetUpsert(true))
if err != nil {
return fmt.Errorf("upsert mcp server: %w", err)
}
return nil
}
func (s *MongoDBStore) Update(ctx context.Context, server ManagedServer) error {
stampUpsert(&server)
result, err := s.collection.UpdateOne(ctx, mongoMCPServerIDFilter{ID: strings.TrimSpace(server.Name)}, bson.M{"$set": mongoMCPServerFields(server)})
if err != nil {
return fmt.Errorf("update mcp server: %w", err)
}
if result.MatchedCount == 0 {
return ErrNotFound
}
return nil
}
// mongoMCPServerFields is the $set document for one server's mutable fields.
func mongoMCPServerFields(server ManagedServer) bson.M {
return bson.M{
"display_name": server.DisplayName,
"url": server.URL,
"transport": server.Transport,
"headers": server.Headers,
"description": server.Description,
"enabled": server.Enabled,
"allowed_tools": server.AllowedTools,
"disallowed_tools": server.DisallowedTools,
"user_paths": server.UserPaths,
"disallowed_user_paths": server.DisallowedUserPaths,
"tool_timeout_seconds": server.ToolTimeoutSeconds,
"updated_at": server.UpdatedAt,
}
}
func (s *MongoDBStore) Delete(ctx context.Context, name string) error {
result, err := s.collection.DeleteOne(ctx, mongoMCPServerIDFilter{ID: strings.TrimSpace(name)})
if err != nil {
return fmt.Errorf("delete mcp server: %w", err)
}
if result.DeletedCount == 0 {
return ErrNotFound
}
return nil
}
// swapHeaders replaces a server's header values only while they still hold
// the values in current, and reports whether it did. Headers added since are
// kept.
func (s *MongoDBStore) swapHeaders(ctx context.Context, current, next ManagedServer) (bool, error) {
if len(current.Headers) == 0 {
return false, nil
}
match, update := storage.MongoSwapSubfields("headers", current.Headers, next.Headers)
match["_id"] = strings.TrimSpace(current.Name)
result, err := s.collection.UpdateOne(ctx, match, update)
if err != nil {
return false, fmt.Errorf("swap mcp server headers: %w", err)
}
return result.MatchedCount > 0, nil
}
func (s *MongoDBStore) Close() error {
return nil
}
func managedServerFromMongo(doc mongoMCPServerDocument) ManagedServer {
server := ManagedServer{
Name: doc.ID,
DisplayName: doc.DisplayName,
URL: doc.URL,
Transport: doc.Transport,
Description: doc.Description,
Enabled: doc.Enabled,
ToolTimeoutSeconds: doc.ToolTimeoutSeconds,
CreatedAt: doc.CreatedAt.UTC(),
UpdatedAt: doc.UpdatedAt.UTC(),
}
if server.DisplayName == "" {
server.DisplayName = server.Name
}
if len(doc.Headers) > 0 {
server.Headers = doc.Headers
}
if len(doc.AllowedTools) > 0 {
server.AllowedTools = append([]string(nil), doc.AllowedTools...)
}
if len(doc.DisallowedTools) > 0 {
server.DisallowedTools = append([]string(nil), doc.DisallowedTools...)
}
if len(doc.UserPaths) > 0 {
server.UserPaths = append([]string(nil), doc.UserPaths...)
}
if len(doc.DisallowedUserPaths) > 0 {
server.DisallowedUserPaths = append([]string(nil), doc.DisallowedUserPaths...)
}
return server
}