Refactor database implementation to use PostgreSQL instead of MongoDB (#289)
Fixes #228 ## Motivation and Context See #19 ## How Has This Been Tested? - Run locally - Run in docker compose - Run infra in minikube - Unit and integration tests all passing :) ## Breaking Changes Yes: users will not be able to access any data they used to have in custom mongodb. We are accepting this breaking change, given the registry is in experimental development status. ## Types of changes <!-- What types of changes does your code introduce? Put an `x` in all the boxes that apply: --> - [ ] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [x] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update ## Checklist <!-- Go over all the following points, and put an `x` in all the boxes that apply. --> - [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 I have not implemented connection pooling support for simplicity, and because cloudnative-pg already provides built-in connection pooling with pgBouncer so if we want it later we can switch to this endpoint easily.
This commit is contained in:
adam jones
GitHub
parent
ee910ca286
commit
ebaf4c9c99
5 changed files with 846 additions and 564 deletions
|
|
@ -36,8 +36,8 @@ type ConnectionType string
|
|||
const (
|
||||
// ConnectionTypeMemory represents an in-memory database connection
|
||||
ConnectionTypeMemory ConnectionType = "memory"
|
||||
// ConnectionTypeMongoDB represents a MongoDB database connection
|
||||
ConnectionTypeMongoDB ConnectionType = "mongodb"
|
||||
// ConnectionTypePostgreSQL represents a PostgreSQL database connection
|
||||
ConnectionTypePostgreSQL ConnectionType = "postgresql"
|
||||
)
|
||||
|
||||
// ConnectionInfo provides information about the database connection
|
||||
|
|
@ -47,7 +47,7 @@ type ConnectionInfo struct {
|
|||
// IsConnected indicates whether the database is currently connected
|
||||
IsConnected bool
|
||||
// Raw provides access to the underlying connection object, which will vary by implementation
|
||||
// For MongoDB, this will be *mongo.Client
|
||||
// For PostgreSQL, this will be *pgx.Conn
|
||||
// For MemoryDB, this will be map[string]*model.MCPRegistry
|
||||
Raw any
|
||||
}
|
||||
|
|
|
|||
192
internal/database/migrate.go
Normal file
192
internal/database/migrate.go
Normal file
|
|
@ -0,0 +1,192 @@
|
|||
package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
)
|
||||
|
||||
//go:embed migrations/*.sql
|
||||
var migrationFiles embed.FS
|
||||
|
||||
// Migration represents a database migration
|
||||
type Migration struct {
|
||||
Version int
|
||||
Name string
|
||||
SQL string
|
||||
}
|
||||
|
||||
// Migrator handles database migrations
|
||||
type Migrator struct {
|
||||
conn *pgx.Conn
|
||||
}
|
||||
|
||||
// NewMigrator creates a new migrator instance
|
||||
func NewMigrator(conn *pgx.Conn) *Migrator {
|
||||
return &Migrator{conn: conn}
|
||||
}
|
||||
|
||||
// ensureMigrationsTable creates the migrations tracking table if it doesn't exist
|
||||
func (m *Migrator) ensureMigrationsTable(ctx context.Context) error {
|
||||
query := `
|
||||
CREATE TABLE IF NOT EXISTS schema_migrations (
|
||||
version INTEGER PRIMARY KEY,
|
||||
name VARCHAR(255) NOT NULL,
|
||||
applied_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
||||
)
|
||||
`
|
||||
_, err := m.conn.Exec(ctx, query)
|
||||
return err
|
||||
}
|
||||
|
||||
// getAppliedMigrations returns a map of already applied migration versions
|
||||
func (m *Migrator) getAppliedMigrations(ctx context.Context) (map[int]bool, error) {
|
||||
query := "SELECT version FROM schema_migrations ORDER BY version"
|
||||
rows, err := m.conn.Query(ctx, query)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to query applied migrations: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
applied := make(map[int]bool)
|
||||
for rows.Next() {
|
||||
var version int
|
||||
if err := rows.Scan(&version); err != nil {
|
||||
return nil, fmt.Errorf("failed to scan migration version: %w", err)
|
||||
}
|
||||
applied[version] = true
|
||||
}
|
||||
|
||||
return applied, rows.Err()
|
||||
}
|
||||
|
||||
// loadMigrations loads all migration files from the embedded filesystem
|
||||
func (m *Migrator) loadMigrations() ([]Migration, error) {
|
||||
entries, err := migrationFiles.ReadDir("migrations")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read migrations directory: %w", err)
|
||||
}
|
||||
|
||||
var migrations []Migration
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sql") {
|
||||
continue
|
||||
}
|
||||
|
||||
// Parse version from filename (e.g., "001_initial_schema.sql" -> version 1)
|
||||
name := entry.Name()
|
||||
parts := strings.SplitN(name, "_", 2)
|
||||
if len(parts) != 2 {
|
||||
log.Printf("Skipping migration file with invalid name format: %s", name)
|
||||
continue
|
||||
}
|
||||
|
||||
version, err := strconv.Atoi(parts[0])
|
||||
if err != nil {
|
||||
log.Printf("Skipping migration file with invalid version: %s", name)
|
||||
continue
|
||||
}
|
||||
|
||||
// Read the migration SQL
|
||||
content, err := migrationFiles.ReadFile(filepath.Join("migrations", name))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read migration file %s: %w", name, err)
|
||||
}
|
||||
|
||||
migrations = append(migrations, Migration{
|
||||
Version: version,
|
||||
Name: strings.TrimSuffix(name, ".sql"),
|
||||
SQL: string(content),
|
||||
})
|
||||
}
|
||||
|
||||
// Sort migrations by version
|
||||
sort.Slice(migrations, func(i, j int) bool {
|
||||
return migrations[i].Version < migrations[j].Version
|
||||
})
|
||||
|
||||
return migrations, nil
|
||||
}
|
||||
|
||||
// Migrate runs all pending migrations
|
||||
func (m *Migrator) Migrate(ctx context.Context) error {
|
||||
// Ensure the migrations table exists
|
||||
if err := m.ensureMigrationsTable(ctx); err != nil {
|
||||
return fmt.Errorf("failed to create migrations table: %w", err)
|
||||
}
|
||||
|
||||
// Get applied migrations
|
||||
applied, err := m.getAppliedMigrations(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to get applied migrations: %w", err)
|
||||
}
|
||||
|
||||
// Load all migration files
|
||||
migrations, err := m.loadMigrations()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load migrations: %w", err)
|
||||
}
|
||||
|
||||
// Find pending migrations
|
||||
var pending []Migration
|
||||
for _, migration := range migrations {
|
||||
if !applied[migration.Version] {
|
||||
pending = append(pending, migration)
|
||||
}
|
||||
}
|
||||
|
||||
if len(pending) == 0 {
|
||||
log.Println("No pending migrations")
|
||||
return nil
|
||||
}
|
||||
|
||||
log.Printf("Applying %d pending migrations", len(pending))
|
||||
|
||||
// Apply each pending migration in a transaction
|
||||
for _, migration := range pending {
|
||||
if err := m.applyMigration(ctx, migration); err != nil {
|
||||
return fmt.Errorf("failed to apply migration %d (%s): %w", migration.Version, migration.Name, err)
|
||||
}
|
||||
log.Printf("Applied migration %d: %s", migration.Version, migration.Name)
|
||||
}
|
||||
|
||||
log.Println("All migrations applied successfully")
|
||||
return nil
|
||||
}
|
||||
|
||||
// applyMigration applies a single migration in a transaction
|
||||
func (m *Migrator) applyMigration(ctx context.Context, migration Migration) error {
|
||||
tx, err := m.conn.Begin(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to begin transaction: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := tx.Rollback(ctx); err != nil && !errors.Is(err, pgx.ErrTxClosed) {
|
||||
log.Printf("Failed to rollback migration transaction: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Execute the migration SQL
|
||||
_, err = tx.Exec(ctx, migration.SQL)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to execute migration SQL: %w", err)
|
||||
}
|
||||
|
||||
// Record the migration as applied
|
||||
_, err = tx.Exec(ctx,
|
||||
"INSERT INTO schema_migrations (version, name) VALUES ($1, $2)",
|
||||
migration.Version, migration.Name)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to record migration: %w", err)
|
||||
}
|
||||
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
58
internal/database/migrations/001_initial_schema.sql
Normal file
58
internal/database/migrations/001_initial_schema.sql
Normal file
|
|
@ -0,0 +1,58 @@
|
|||
-- PostgreSQL schema for MCP Registry
|
||||
-- This schema uses a hybrid approach with relational tables for core entities
|
||||
-- and JSONB columns for complex nested data structures
|
||||
|
||||
-- Enable UUID extension
|
||||
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
|
||||
CREATE EXTENSION IF NOT EXISTS "pg_trgm";
|
||||
|
||||
-- Servers table - core server information
|
||||
CREATE TABLE servers (
|
||||
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
|
||||
name VARCHAR(255) NOT NULL,
|
||||
description TEXT NOT NULL,
|
||||
status VARCHAR(50) DEFAULT 'active',
|
||||
|
||||
-- Repository information stored as JSONB for flexibility
|
||||
repository JSONB,
|
||||
|
||||
-- Version details
|
||||
version VARCHAR(255) NOT NULL,
|
||||
release_date TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||||
is_latest BOOLEAN DEFAULT true,
|
||||
|
||||
-- Packages and remotes stored as JSONB arrays
|
||||
packages JSONB DEFAULT '[]'::jsonb,
|
||||
remotes JSONB DEFAULT '[]'::jsonb,
|
||||
|
||||
-- Metadata
|
||||
created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
|
||||
updated_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- Indexes for performance
|
||||
CREATE UNIQUE INDEX idx_servers_id ON servers(id);
|
||||
CREATE INDEX idx_servers_name ON servers(name);
|
||||
CREATE UNIQUE INDEX idx_servers_name_version ON servers(name, version);
|
||||
CREATE INDEX idx_servers_latest ON servers(is_latest) WHERE is_latest = true;
|
||||
CREATE INDEX idx_servers_status ON servers(status);
|
||||
|
||||
-- GIN indexes for JSONB fields
|
||||
CREATE INDEX idx_servers_repository_gin ON servers USING GIN(repository);
|
||||
CREATE INDEX idx_servers_packages_gin ON servers USING GIN(packages);
|
||||
CREATE INDEX idx_servers_remotes_gin ON servers USING GIN(remotes);
|
||||
|
||||
-- Function to update the updated_at timestamp
|
||||
CREATE OR REPLACE FUNCTION update_updated_at_column()
|
||||
RETURNS TRIGGER AS $$
|
||||
BEGIN
|
||||
NEW.updated_at = NOW();
|
||||
RETURN NEW;
|
||||
END;
|
||||
$$ language 'plpgsql';
|
||||
|
||||
-- Trigger to automatically update updated_at
|
||||
CREATE TRIGGER update_servers_updated_at
|
||||
BEFORE UPDATE ON servers
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION update_updated_at_column();
|
||||
|
|
@ -1,309 +0,0 @@
|
|||
package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"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"
|
||||
)
|
||||
|
||||
// MongoDB is an implementation of the Database interface using MongoDB
|
||||
type MongoDB struct {
|
||||
client *mongo.Client
|
||||
database *mongo.Database
|
||||
collection *mongo.Collection
|
||||
}
|
||||
|
||||
// NewMongoDB creates a new instance of the MongoDB database
|
||||
func NewMongoDB(ctx context.Context, connectionURI, databaseName, collectionName string) (*MongoDB, error) {
|
||||
// Set client options and connect to MongoDB
|
||||
clientOptions := options.Client().ApplyURI(connectionURI)
|
||||
client, err := mongo.Connect(ctx, clientOptions)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Ping the MongoDB server to verify the connection
|
||||
if err = client.Ping(ctx, nil); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Get database and collection
|
||||
database := client.Database(databaseName)
|
||||
collection := database.Collection(collectionName)
|
||||
|
||||
// Create indexes for better query performance
|
||||
models := []mongo.IndexModel{
|
||||
{
|
||||
Keys: bson.D{bson.E{Key: "name", Value: 1}},
|
||||
},
|
||||
{
|
||||
Keys: bson.D{bson.E{Key: "id", Value: 1}},
|
||||
Options: options.Index().SetUnique(true),
|
||||
},
|
||||
// add an index for the combination of name and version
|
||||
{
|
||||
Keys: bson.D{bson.E{Key: "name", Value: 1}, bson.E{Key: "version_detail.version", Value: 1}},
|
||||
Options: options.Index().SetUnique(true),
|
||||
},
|
||||
}
|
||||
|
||||
_, err = collection.Indexes().CreateMany(ctx, models)
|
||||
if err != nil {
|
||||
// Mongo will error if the index already exists, we can ignore this and continue.
|
||||
var commandError mongo.CommandError
|
||||
if errors.As(err, &commandError) && commandError.Code != 86 {
|
||||
return nil, err
|
||||
}
|
||||
log.Printf("Indexes already exists, skipping.")
|
||||
}
|
||||
|
||||
return &MongoDB{
|
||||
client: client,
|
||||
database: database,
|
||||
collection: collection,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// List retrieves MCPRegistry entries with optional filtering and pagination
|
||||
func (db *MongoDB) List(
|
||||
ctx context.Context,
|
||||
filter map[string]any,
|
||||
cursor string,
|
||||
limit int,
|
||||
) ([]*model.Server, string, error) {
|
||||
if limit <= 0 {
|
||||
// Set default limit if not provided
|
||||
limit = 10
|
||||
}
|
||||
|
||||
if ctx.Err() != nil {
|
||||
return nil, "", ctx.Err()
|
||||
}
|
||||
|
||||
// Convert Go map to MongoDB filter
|
||||
mongoFilter := bson.M{
|
||||
"version_detail.is_latest": true,
|
||||
}
|
||||
// Map common filter keys to MongoDB document paths
|
||||
for k, v := range filter {
|
||||
// Handle nested fields with dot notation
|
||||
switch k {
|
||||
case "version":
|
||||
mongoFilter["version_detail.version"] = v
|
||||
case "name":
|
||||
mongoFilter["name"] = v
|
||||
default:
|
||||
mongoFilter[k] = v
|
||||
}
|
||||
}
|
||||
|
||||
// Setup pagination options
|
||||
findOptions := options.Find()
|
||||
|
||||
// If cursor is provided, add condition to filter to only get records after the cursor
|
||||
if cursor != "" {
|
||||
// Validate that the cursor is a valid UUID
|
||||
if _, err := uuid.Parse(cursor); err != nil {
|
||||
return nil, "", fmt.Errorf("invalid cursor format: %w", err)
|
||||
}
|
||||
|
||||
// Fetch the document at the cursor to get its sort values
|
||||
var cursorDoc model.Server
|
||||
err := db.collection.FindOne(ctx, bson.M{"id": cursor}).Decode(&cursorDoc)
|
||||
if err != nil {
|
||||
if !errors.Is(err, mongo.ErrNoDocuments) {
|
||||
return nil, "", err
|
||||
}
|
||||
// If cursor document not found, start from beginning
|
||||
} else {
|
||||
// Use the cursor document's ID to paginate (records with ID > cursor's ID)
|
||||
mongoFilter["id"] = bson.M{"$gt": cursor}
|
||||
}
|
||||
}
|
||||
|
||||
// Set sort order by ID (for consistent pagination)
|
||||
findOptions.SetSort(bson.M{"id": 1})
|
||||
|
||||
// Set limit if provided and valid
|
||||
if limit > 0 {
|
||||
findOptions.SetLimit(int64(limit))
|
||||
}
|
||||
|
||||
// Execute find operation with options
|
||||
mongoCursor, err := db.collection.Find(ctx, mongoFilter, findOptions)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
defer mongoCursor.Close(ctx)
|
||||
|
||||
// Decode results
|
||||
var results []*model.Server
|
||||
if err = mongoCursor.All(ctx, &results); err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
// Determine the next cursor
|
||||
nextCursor := ""
|
||||
if len(results) > 0 && limit > 0 && len(results) >= limit {
|
||||
// Use the last item's ID as the next cursor
|
||||
nextCursor = results[len(results)-1].ID
|
||||
}
|
||||
|
||||
return results, nextCursor, nil
|
||||
}
|
||||
|
||||
// GetByID retrieves a single ServerDetail by its ID
|
||||
func (db *MongoDB) GetByID(ctx context.Context, id string) (*model.ServerDetail, error) {
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
// Create a filter for the ID
|
||||
filter := bson.M{"id": id}
|
||||
|
||||
// Find the entry in the database
|
||||
var entry model.ServerDetail
|
||||
err := db.collection.FindOne(ctx, filter).Decode(&entry)
|
||||
if err != nil {
|
||||
if errors.Is(err, mongo.ErrNoDocuments) {
|
||||
return nil, ErrNotFound
|
||||
}
|
||||
return nil, fmt.Errorf("error retrieving entry: %w", err)
|
||||
}
|
||||
|
||||
// Create and return a ServerDetail from the entry data
|
||||
return &entry, nil
|
||||
}
|
||||
|
||||
// Publish adds a new ServerDetail to the database
|
||||
func (db *MongoDB) Publish(ctx context.Context, serverDetail *model.ServerDetail) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
// find a server detail with the same name and check that the current version is greater than the existing one
|
||||
filter := bson.M{
|
||||
"name": serverDetail.Name,
|
||||
"version_detail.is_latest": true,
|
||||
}
|
||||
|
||||
var existingEntry model.ServerDetail
|
||||
err := db.collection.FindOne(ctx, filter).Decode(&existingEntry)
|
||||
if err != nil && !errors.Is(err, mongo.ErrNoDocuments) {
|
||||
return fmt.Errorf("error checking existing entry: %w", err)
|
||||
}
|
||||
|
||||
// check that the current version is greater than the existing one
|
||||
if serverDetail.VersionDetail.Version <= existingEntry.VersionDetail.Version {
|
||||
return fmt.Errorf("version must be greater than existing version")
|
||||
}
|
||||
|
||||
serverDetail.ID = uuid.New().String()
|
||||
serverDetail.VersionDetail.IsLatest = true
|
||||
serverDetail.VersionDetail.ReleaseDate = time.Now().Format(time.RFC3339)
|
||||
|
||||
// Insert the entry into the database
|
||||
_, err = db.collection.InsertOne(ctx, serverDetail)
|
||||
if err != nil {
|
||||
if mongo.IsDuplicateKeyError(err) {
|
||||
return ErrAlreadyExists
|
||||
}
|
||||
return fmt.Errorf("error inserting entry: %w", err)
|
||||
}
|
||||
|
||||
// update the existing entry to not be the latest version
|
||||
if existingEntry.ID != "" {
|
||||
_, err = db.collection.UpdateOne(
|
||||
ctx,
|
||||
bson.M{"id": existingEntry.ID},
|
||||
bson.M{"$set": bson.M{"version_detail.islatest": false}})
|
||||
if err != nil {
|
||||
return fmt.Errorf("error updating existing entry: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ImportSeed imports initial data from a seed file into MongoDB
|
||||
func (db *MongoDB) ImportSeed(ctx context.Context, seedFilePath string) error {
|
||||
// Read the seed file
|
||||
servers, err := ReadSeedFile(ctx, seedFilePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read seed file: %w", err)
|
||||
}
|
||||
|
||||
collection := db.collection
|
||||
|
||||
log.Printf("Importing %d servers into collection %s", len(servers), collection.Name())
|
||||
|
||||
for i, server := range servers {
|
||||
if server.ID == "" || server.Name == "" {
|
||||
log.Printf("Skipping server %d: ID or Name is empty", i+1)
|
||||
continue
|
||||
}
|
||||
|
||||
if server.VersionDetail.Version == "" {
|
||||
server.VersionDetail.Version = "0.0.1-seed"
|
||||
server.VersionDetail.ReleaseDate = time.Now().Format(time.RFC3339)
|
||||
server.VersionDetail.IsLatest = true
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
switch {
|
||||
case result.UpsertedCount > 0:
|
||||
log.Printf("[%d/%d] Created server: %s", i+1, len(servers), server.Name)
|
||||
case result.ModifiedCount > 0:
|
||||
log.Printf("[%d/%d] Updated server: %s", i+1, len(servers), server.Name)
|
||||
default:
|
||||
log.Printf("[%d/%d] Server already up to date: %s", i+1, len(servers), server.Name)
|
||||
}
|
||||
}
|
||||
|
||||
log.Println("MongoDB database import completed successfully")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close closes the database connection
|
||||
func (db *MongoDB) Close() error {
|
||||
return db.client.Disconnect(context.Background())
|
||||
}
|
||||
|
||||
// Connection returns information about the database connection
|
||||
func (db *MongoDB) Connection() *ConnectionInfo {
|
||||
isConnected := false
|
||||
// Check if the client is connected
|
||||
if db.client != nil {
|
||||
// A quick ping with 1 second timeout to verify connection
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
err := db.client.Ping(ctx, nil)
|
||||
isConnected = (err == nil)
|
||||
}
|
||||
|
||||
return &ConnectionInfo{
|
||||
Type: ConnectionTypeMongoDB,
|
||||
IsConnected: isConnected,
|
||||
Raw: db.client,
|
||||
}
|
||||
}
|
||||
421
internal/database/postgres.go
Normal file
421
internal/database/postgres.go
Normal file
|
|
@ -0,0 +1,421 @@
|
|||
package database
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/modelcontextprotocol/registry/internal/model"
|
||||
)
|
||||
|
||||
// PostgreSQL is an implementation of the Database interface using PostgreSQL
|
||||
type PostgreSQL struct {
|
||||
conn *pgx.Conn
|
||||
}
|
||||
|
||||
// NewPostgreSQL creates a new instance of the PostgreSQL database
|
||||
func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, error) {
|
||||
conn, err := pgx.Connect(ctx, connectionURI)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to connect to PostgreSQL: %w", err)
|
||||
}
|
||||
|
||||
// Test the connection
|
||||
if err = conn.Ping(ctx); err != nil {
|
||||
return nil, fmt.Errorf("failed to ping PostgreSQL: %w", err)
|
||||
}
|
||||
|
||||
// Run migrations
|
||||
migrator := NewMigrator(conn)
|
||||
if err := migrator.Migrate(ctx); err != nil {
|
||||
return nil, fmt.Errorf("failed to run database migrations: %w", err)
|
||||
}
|
||||
|
||||
return &PostgreSQL{
|
||||
conn: conn,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// List retrieves MCPRegistry entries with optional filtering and pagination
|
||||
func (db *PostgreSQL) List(
|
||||
ctx context.Context,
|
||||
filter map[string]any,
|
||||
cursor string,
|
||||
limit int,
|
||||
) ([]*model.Server, string, error) {
|
||||
if limit <= 0 {
|
||||
limit = 10
|
||||
}
|
||||
|
||||
if ctx.Err() != nil {
|
||||
return nil, "", ctx.Err()
|
||||
}
|
||||
|
||||
// Build WHERE clause
|
||||
whereClause := "WHERE is_latest = true"
|
||||
args := []any{}
|
||||
argIndex := 1
|
||||
|
||||
// Add filters
|
||||
for k, v := range filter {
|
||||
switch k {
|
||||
case "name":
|
||||
whereClause += fmt.Sprintf(" AND name = $%d", argIndex)
|
||||
args = append(args, v)
|
||||
argIndex++
|
||||
case "version":
|
||||
whereClause += fmt.Sprintf(" AND version = $%d", argIndex)
|
||||
args = append(args, v)
|
||||
argIndex++
|
||||
case "status":
|
||||
whereClause += fmt.Sprintf(" AND status = $%d", argIndex)
|
||||
args = append(args, v)
|
||||
argIndex++
|
||||
}
|
||||
}
|
||||
|
||||
// Add cursor pagination
|
||||
if cursor != "" {
|
||||
if _, err := uuid.Parse(cursor); err != nil {
|
||||
return nil, "", fmt.Errorf("invalid cursor format: %w", err)
|
||||
}
|
||||
whereClause += fmt.Sprintf(" AND id > $%d", argIndex)
|
||||
args = append(args, cursor)
|
||||
argIndex++
|
||||
}
|
||||
|
||||
// Build query
|
||||
query := fmt.Sprintf(`
|
||||
SELECT id, name, description, status, repository, version, release_date, is_latest
|
||||
FROM servers
|
||||
%s
|
||||
ORDER BY id
|
||||
LIMIT $%d
|
||||
`, whereClause, argIndex)
|
||||
args = append(args, limit)
|
||||
|
||||
rows, err := db.conn.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("failed to query servers: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var results []*model.Server
|
||||
for rows.Next() {
|
||||
var server model.Server
|
||||
var repositoryJSON []byte
|
||||
var releaseDate time.Time
|
||||
|
||||
err := rows.Scan(
|
||||
&server.ID,
|
||||
&server.Name,
|
||||
&server.Description,
|
||||
&server.Status,
|
||||
&repositoryJSON,
|
||||
&server.VersionDetail.Version,
|
||||
&releaseDate,
|
||||
&server.VersionDetail.IsLatest,
|
||||
)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("failed to scan server row: %w", err)
|
||||
}
|
||||
|
||||
// Parse repository JSON
|
||||
if len(repositoryJSON) > 0 {
|
||||
if err := json.Unmarshal(repositoryJSON, &server.Repository); err != nil {
|
||||
return nil, "", fmt.Errorf("failed to unmarshal repository: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
server.VersionDetail.ReleaseDate = releaseDate.Format(time.RFC3339)
|
||||
results = append(results, &server)
|
||||
}
|
||||
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, "", fmt.Errorf("error iterating rows: %w", err)
|
||||
}
|
||||
|
||||
// Determine next cursor
|
||||
nextCursor := ""
|
||||
if len(results) > 0 && len(results) >= limit {
|
||||
nextCursor = results[len(results)-1].ID
|
||||
}
|
||||
|
||||
return results, nextCursor, nil
|
||||
}
|
||||
|
||||
// GetByID retrieves a single ServerDetail by its ID
|
||||
func (db *PostgreSQL) GetByID(ctx context.Context, id string) (*model.ServerDetail, error) {
|
||||
if ctx.Err() != nil {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
|
||||
query := `
|
||||
SELECT id, name, description, status, repository, version, release_date, is_latest, packages, remotes
|
||||
FROM servers
|
||||
WHERE id = $1
|
||||
`
|
||||
|
||||
var serverDetail model.ServerDetail
|
||||
var repositoryJSON, packagesJSON, remotesJSON []byte
|
||||
var releaseDate time.Time
|
||||
|
||||
err := db.conn.QueryRow(ctx, query, id).Scan(
|
||||
&serverDetail.ID,
|
||||
&serverDetail.Name,
|
||||
&serverDetail.Description,
|
||||
&serverDetail.Status,
|
||||
&repositoryJSON,
|
||||
&serverDetail.VersionDetail.Version,
|
||||
&releaseDate,
|
||||
&serverDetail.VersionDetail.IsLatest,
|
||||
&packagesJSON,
|
||||
&remotesJSON,
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
return nil, ErrNotFound
|
||||
}
|
||||
return nil, fmt.Errorf("failed to get server by ID: %w", err)
|
||||
}
|
||||
|
||||
// Parse JSON fields
|
||||
if len(repositoryJSON) > 0 {
|
||||
if err := json.Unmarshal(repositoryJSON, &serverDetail.Repository); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal repository: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if len(packagesJSON) > 0 {
|
||||
if err := json.Unmarshal(packagesJSON, &serverDetail.Packages); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal packages: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
if len(remotesJSON) > 0 {
|
||||
if err := json.Unmarshal(remotesJSON, &serverDetail.Remotes); err != nil {
|
||||
return nil, fmt.Errorf("failed to unmarshal remotes: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
serverDetail.VersionDetail.ReleaseDate = releaseDate.Format(time.RFC3339)
|
||||
|
||||
return &serverDetail, nil
|
||||
}
|
||||
|
||||
// Publish adds a new ServerDetail to the database
|
||||
func (db *PostgreSQL) Publish(ctx context.Context, serverDetail *model.ServerDetail) error {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
|
||||
tx, err := db.conn.Begin(ctx)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to begin transaction: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
if err := tx.Rollback(ctx); err != nil {
|
||||
log.Printf("Failed to rollback transaction: %v", err)
|
||||
}
|
||||
}()
|
||||
|
||||
// Check if there's an existing latest version for this server
|
||||
var existingVersion string
|
||||
checkQuery := `
|
||||
SELECT version
|
||||
FROM servers
|
||||
WHERE name = $1 AND is_latest = true
|
||||
`
|
||||
err = tx.QueryRow(ctx, checkQuery, serverDetail.Name).Scan(&existingVersion)
|
||||
if err != nil && !errors.Is(err, pgx.ErrNoRows) {
|
||||
return fmt.Errorf("failed to check existing version: %w", err)
|
||||
}
|
||||
|
||||
// Validate version ordering
|
||||
if existingVersion != "" && serverDetail.VersionDetail.Version <= existingVersion {
|
||||
return ErrInvalidVersion
|
||||
}
|
||||
|
||||
// Prepare JSON data
|
||||
repositoryJSON, err := json.Marshal(serverDetail.Repository)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal repository: %w", err)
|
||||
}
|
||||
|
||||
packagesJSON, err := json.Marshal(serverDetail.Packages)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal packages: %w", err)
|
||||
}
|
||||
|
||||
remotesJSON, err := json.Marshal(serverDetail.Remotes)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to marshal remotes: %w", err)
|
||||
}
|
||||
|
||||
// Generate ID and set metadata
|
||||
if serverDetail.ID == "" {
|
||||
serverDetail.ID = uuid.New().String()
|
||||
}
|
||||
serverDetail.VersionDetail.IsLatest = true
|
||||
if serverDetail.VersionDetail.ReleaseDate == "" {
|
||||
serverDetail.VersionDetail.ReleaseDate = time.Now().Format(time.RFC3339)
|
||||
}
|
||||
|
||||
releaseDate, err := time.Parse(time.RFC3339, serverDetail.VersionDetail.ReleaseDate)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to parse release date: %w", err)
|
||||
}
|
||||
|
||||
// Update existing latest version to not be latest
|
||||
if existingVersion != "" {
|
||||
updateQuery := `
|
||||
UPDATE servers
|
||||
SET is_latest = false
|
||||
WHERE name = $1 AND is_latest = true
|
||||
`
|
||||
_, err = tx.Exec(ctx, updateQuery, serverDetail.Name)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to update existing latest version: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Insert new server version
|
||||
insertQuery := `
|
||||
INSERT INTO servers (id, name, description, status, repository, version, release_date, is_latest, packages, remotes)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
`
|
||||
_, err = tx.Exec(ctx, insertQuery,
|
||||
serverDetail.ID,
|
||||
serverDetail.Name,
|
||||
serverDetail.Description,
|
||||
serverDetail.Status,
|
||||
repositoryJSON,
|
||||
serverDetail.VersionDetail.Version,
|
||||
releaseDate,
|
||||
serverDetail.VersionDetail.IsLatest,
|
||||
packagesJSON,
|
||||
remotesJSON,
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to insert server: %w", err)
|
||||
}
|
||||
|
||||
return tx.Commit(ctx)
|
||||
}
|
||||
|
||||
// ImportSeed imports initial data from a seed file into PostgreSQL
|
||||
func (db *PostgreSQL) ImportSeed(ctx context.Context, seedFilePath string) error {
|
||||
servers, err := ReadSeedFile(ctx, seedFilePath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to read seed file: %w", err)
|
||||
}
|
||||
|
||||
log.Printf("Importing %d servers into PostgreSQL", len(servers))
|
||||
|
||||
for i, server := range servers {
|
||||
if server.ID == "" || server.Name == "" {
|
||||
log.Printf("Skipping server %d: ID or Name is empty", i+1)
|
||||
continue
|
||||
}
|
||||
|
||||
if server.VersionDetail.Version == "" {
|
||||
server.VersionDetail.Version = "0.0.1-seed"
|
||||
server.VersionDetail.ReleaseDate = time.Now().Format(time.RFC3339)
|
||||
server.VersionDetail.IsLatest = true
|
||||
}
|
||||
|
||||
// Prepare JSON data
|
||||
repositoryJSON, err := json.Marshal(server.Repository)
|
||||
if err != nil {
|
||||
log.Printf("Error marshaling repository for server %s: %v", server.ID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
packagesJSON, err := json.Marshal(server.Packages)
|
||||
if err != nil {
|
||||
log.Printf("Error marshaling packages for server %s: %v", server.ID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
remotesJSON, err := json.Marshal(server.Remotes)
|
||||
if err != nil {
|
||||
log.Printf("Error marshaling remotes for server %s: %v", server.ID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
releaseDate, err := time.Parse(time.RFC3339, server.VersionDetail.ReleaseDate)
|
||||
if err != nil {
|
||||
log.Printf("Error parsing release date for server %s: %v", server.ID, err)
|
||||
releaseDate = time.Now()
|
||||
}
|
||||
|
||||
// Use UPSERT (INSERT ... ON CONFLICT DO UPDATE)
|
||||
upsertQuery := `
|
||||
INSERT INTO servers (id, name, description, status, repository, version, release_date, is_latest, packages, remotes)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
ON CONFLICT (id) DO UPDATE SET
|
||||
name = EXCLUDED.name,
|
||||
description = EXCLUDED.description,
|
||||
status = EXCLUDED.status,
|
||||
repository = EXCLUDED.repository,
|
||||
version = EXCLUDED.version,
|
||||
release_date = EXCLUDED.release_date,
|
||||
is_latest = EXCLUDED.is_latest,
|
||||
packages = EXCLUDED.packages,
|
||||
remotes = EXCLUDED.remotes,
|
||||
updated_at = NOW()
|
||||
`
|
||||
|
||||
_, err = db.conn.Exec(ctx, upsertQuery,
|
||||
server.ID,
|
||||
server.Name,
|
||||
server.Description,
|
||||
server.Status,
|
||||
repositoryJSON,
|
||||
server.VersionDetail.Version,
|
||||
releaseDate,
|
||||
server.VersionDetail.IsLatest,
|
||||
packagesJSON,
|
||||
remotesJSON,
|
||||
)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("Error importing server %s: %v", server.ID, err)
|
||||
continue
|
||||
}
|
||||
|
||||
log.Printf("[%d/%d] Imported server: %s", i+1, len(servers), server.Name)
|
||||
}
|
||||
|
||||
log.Println("PostgreSQL database import completed successfully")
|
||||
return nil
|
||||
}
|
||||
|
||||
// Close closes the database connection
|
||||
func (db *PostgreSQL) Close() error {
|
||||
return db.conn.Close(context.Background())
|
||||
}
|
||||
|
||||
// Connection returns information about the database connection
|
||||
func (db *PostgreSQL) Connection() *ConnectionInfo {
|
||||
isConnected := false
|
||||
if db.conn != nil {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
err := db.conn.Ping(ctx)
|
||||
isConnected = (err == nil)
|
||||
}
|
||||
|
||||
return &ConnectionInfo{
|
||||
Type: ConnectionTypePostgreSQL,
|
||||
IsConnected: isConnected,
|
||||
Raw: db.conn,
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue