From 747ad310bba75d80f20c1dd03051dd825df734e0 Mon Sep 17 00:00:00 2001 From: adam jones Date: Wed, 24 Sep 2025 18:27:36 +0100 Subject: [PATCH] Fix atomic latest version update to prevent missing isLatest flags (#530) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Fixes a race condition where a server could have no version marked as `isLatest` if the database operations failed independently. ## Changes - **CreateServer signature**: Now accepts `oldLatestVersionID` parameter to atomically unmark previous latest version - **PostgreSQL implementation**: Uses a transaction to ensure UPDATE (unmark old latest) and INSERT (create new version) happen atomically - **Service layer**: Maintains `WithPublishLock` for concurrent operation serialization - **Two-level protection**: Advisory lock prevents races, transaction ensures atomicity ## Test plan - [x] Existing race condition tests pass - [x] Integration tests pass - [x] Unit tests pass 🤖 Generated with [Claude Code](https://claude.ai/code) --------- Co-authored-by: Claude --- internal/database/database.go | 34 ++-- internal/database/postgres.go | 104 +++++++----- internal/database/testutil.go | 195 ++++++++++------------ internal/importer/importer.go | 4 +- internal/importer/importer_test.go | 6 +- internal/service/registry_service.go | 110 ++++++------ internal/service/registry_service_test.go | 25 ++- 7 files changed, 244 insertions(+), 234 deletions(-) diff --git a/internal/database/database.go b/internal/database/database.go index f1bd477..d3d076e 100644 --- a/internal/database/database.go +++ b/internal/database/database.go @@ -5,6 +5,7 @@ import ( "errors" "time" + "github.com/jackc/pgx/v5" apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0" ) @@ -31,36 +32,38 @@ type ServerFilter struct { // Database defines the interface for database operations type Database interface { // Retrieve server entries with optional filtering - List(ctx context.Context, filter *ServerFilter, cursor string, limit int) ([]*apiv0.ServerJSON, string, error) + List(ctx context.Context, tx pgx.Tx, filter *ServerFilter, cursor string, limit int) ([]*apiv0.ServerJSON, string, error) // Retrieve a single server by its version ID - GetByVersionID(ctx context.Context, versionID string) (*apiv0.ServerJSON, error) + GetByVersionID(ctx context.Context, tx pgx.Tx, versionID string) (*apiv0.ServerJSON, error) // Retrieve latest version of a server by server ID - GetByServerID(ctx context.Context, serverID string) (*apiv0.ServerJSON, error) + GetByServerID(ctx context.Context, tx pgx.Tx, serverID string) (*apiv0.ServerJSON, error) // Retrieve specific version of a server by server ID and version - GetByServerIDAndVersion(ctx context.Context, serverID string, version string) (*apiv0.ServerJSON, error) + GetByServerIDAndVersion(ctx context.Context, tx pgx.Tx, serverID string, version string) (*apiv0.ServerJSON, error) // Retrieve all versions of a server by server ID - GetAllVersionsByServerID(ctx context.Context, serverID string) ([]*apiv0.ServerJSON, error) - // CreateServer adds a new server to the database - CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) + GetAllVersionsByServerID(ctx context.Context, tx pgx.Tx, serverID string) ([]*apiv0.ServerJSON, error) + // CreateServer inserts a new server version + CreateServer(ctx context.Context, tx pgx.Tx, newServer *apiv0.ServerJSON) (*apiv0.ServerJSON, error) // UpdateServer updates an existing server record - UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) - // WithPublishLock executes a function with an exclusive lock for publishing a server + UpdateServer(ctx context.Context, tx pgx.Tx, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) + // AcquirePublishLock acquires an exclusive advisory lock for publishing a server // This prevents race conditions when multiple versions are published concurrently - WithPublishLock(ctx context.Context, serverName string, fn func(ctx context.Context) error) error + AcquirePublishLock(ctx context.Context, tx pgx.Tx, serverName string) error + // InTransaction executes a function within a database transaction + InTransaction(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) error // Close closes the database connection Close() error } -// WithPublishLockT is a generic helper that wraps WithPublishLock for functions returning a value +// InTransactionT is a generic helper that wraps InTransaction for functions returning a value // This exists because Go does not support generic methods on interfaces - only the Database interface -// method WithPublishLock (without generics) can exist, so we provide this generic wrapper function. +// method InTransaction (without generics) can exist, so we provide this generic wrapper function. // This is a common pattern in Go for working around this language limitation. -func WithPublishLockT[T any](ctx context.Context, db Database, serverName string, fn func(ctx context.Context) (T, error)) (T, error) { +func InTransactionT[T any](ctx context.Context, db Database, fn func(ctx context.Context, tx pgx.Tx) (T, error)) (T, error) { var result T var fnErr error - err := db.WithPublishLock(ctx, serverName, func(lockCtx context.Context) error { - result, fnErr = fn(lockCtx) + err := db.InTransaction(ctx, func(txCtx context.Context, tx pgx.Tx) error { + result, fnErr = fn(txCtx, tx) return fnErr }) @@ -71,4 +74,3 @@ func WithPublishLockT[T any](ctx context.Context, db Database, serverName string return result, nil } - diff --git a/internal/database/postgres.go b/internal/database/postgres.go index 10bfb51..450810b 100644 --- a/internal/database/postgres.go +++ b/internal/database/postgres.go @@ -5,11 +5,13 @@ import ( "encoding/json" "errors" "fmt" + "log" "strings" "time" "github.com/google/uuid" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/jackc/pgx/v5/pgxpool" apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0" @@ -20,6 +22,21 @@ type PostgreSQL struct { pool *pgxpool.Pool } +// Executor is an interface for executing queries (satisfied by both pgx.Tx and pgxpool.Pool) +type Executor interface { + Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error) + Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error) + QueryRow(ctx context.Context, sql string, args ...any) pgx.Row +} + +// getExecutor returns the appropriate executor (transaction or pool) +func (db *PostgreSQL) getExecutor(tx pgx.Tx) Executor { + if tx != nil { + return tx + } + return db.pool +} + // NewPostgreSQL creates a new instance of the PostgreSQL database func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, error) { // Parse connection config for pool settings @@ -65,6 +82,7 @@ func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, erro //nolint:cyclop // Database filtering logic is inherently complex but clear func (db *PostgreSQL) List( ctx context.Context, + tx pgx.Tx, filter *ServerFilter, cursor string, limit int, @@ -142,7 +160,7 @@ func (db *PostgreSQL) List( `, whereClause, argIndex) args = append(args, limit) - rows, err := db.pool.Query(ctx, query, args...) + rows, err := db.getExecutor(tx).Query(ctx, query, args...) if err != nil { return nil, "", fmt.Errorf("failed to query servers: %w", err) } @@ -182,7 +200,7 @@ func (db *PostgreSQL) List( return results, nextCursor, nil } -func (db *PostgreSQL) GetByVersionID(ctx context.Context, versionID string) (*apiv0.ServerJSON, error) { +func (db *PostgreSQL) GetByVersionID(ctx context.Context, tx pgx.Tx, versionID string) (*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } @@ -194,8 +212,7 @@ func (db *PostgreSQL) GetByVersionID(ctx context.Context, versionID string) (*ap ` var valueJSON []byte - - err := db.pool.QueryRow(ctx, query, versionID).Scan(&valueJSON) + err := db.getExecutor(tx).QueryRow(ctx, query, versionID).Scan(&valueJSON) if err != nil { if errors.Is(err, pgx.ErrNoRows) { @@ -214,7 +231,7 @@ func (db *PostgreSQL) GetByVersionID(ctx context.Context, versionID string) (*ap } // GetByServerID retrieves the latest version of a server by server ID -func (db *PostgreSQL) GetByServerID(ctx context.Context, serverID string) (*apiv0.ServerJSON, error) { +func (db *PostgreSQL) GetByServerID(ctx context.Context, tx pgx.Tx, serverID string) (*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } @@ -228,8 +245,7 @@ func (db *PostgreSQL) GetByServerID(ctx context.Context, serverID string) (*apiv ` var valueJSON []byte - - err := db.pool.QueryRow(ctx, query, serverID).Scan(&valueJSON) + err := db.getExecutor(tx).QueryRow(ctx, query, serverID).Scan(&valueJSON) if err != nil { if errors.Is(err, pgx.ErrNoRows) { @@ -248,7 +264,7 @@ func (db *PostgreSQL) GetByServerID(ctx context.Context, serverID string) (*apiv } // GetByServerIDAndVersion retrieves a specific version of a server by server ID and version -func (db *PostgreSQL) GetByServerIDAndVersion(ctx context.Context, serverID string, version string) (*apiv0.ServerJSON, error) { +func (db *PostgreSQL) GetByServerIDAndVersion(ctx context.Context, tx pgx.Tx, serverID string, version string) (*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } @@ -261,8 +277,7 @@ func (db *PostgreSQL) GetByServerIDAndVersion(ctx context.Context, serverID stri ` var valueJSON []byte - - err := db.pool.QueryRow(ctx, query, serverID, version).Scan(&valueJSON) + err := db.getExecutor(tx).QueryRow(ctx, query, serverID, version).Scan(&valueJSON) if err != nil { if errors.Is(err, pgx.ErrNoRows) { @@ -281,7 +296,7 @@ func (db *PostgreSQL) GetByServerIDAndVersion(ctx context.Context, serverID stri } // GetAllVersionsByServerID retrieves all versions of a server by server ID -func (db *PostgreSQL) GetAllVersionsByServerID(ctx context.Context, serverID string) ([]*apiv0.ServerJSON, error) { +func (db *PostgreSQL) GetAllVersionsByServerID(ctx context.Context, tx pgx.Tx, serverID string) ([]*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } @@ -293,7 +308,7 @@ func (db *PostgreSQL) GetAllVersionsByServerID(ctx context.Context, serverID str ORDER BY (value->'_meta'->'io.modelcontextprotocol.registry/official'->>'publishedAt')::timestamp DESC ` - rows, err := db.pool.Query(ctx, query, serverID) + rows, err := db.getExecutor(tx).Query(ctx, query, serverID) if err != nil { return nil, fmt.Errorf("failed to query server versions: %w", err) } @@ -328,8 +343,8 @@ func (db *PostgreSQL) GetAllVersionsByServerID(ctx context.Context, serverID str return results, nil } -// CreateServer adds a new server to the database -func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) { +// CreateServer inserts a new server version +func (db *PostgreSQL) CreateServer(ctx context.Context, tx pgx.Tx, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } @@ -339,11 +354,9 @@ func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON return nil, fmt.Errorf("server must have registry metadata with ServerID and VersionID") } - serverID := server.Meta.Official.ServerID versionID := server.Meta.Official.VersionID - - if serverID == "" || versionID == "" { - return nil, fmt.Errorf("server must have both ServerID and VersionID in registry metadata") + if versionID == "" { + return nil, fmt.Errorf("server must have VersionID in registry metadata") } // Marshal the complete server to JSONB @@ -352,13 +365,14 @@ func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON return nil, fmt.Errorf("failed to marshal server JSON: %w", err) } - // Insert into servers table with new schema (only version_id column, serverId is in JSON) - query := ` + // Insert the new version + insertQuery := ` INSERT INTO servers (version_id, value) VALUES ($1, $2) ` - _, err = db.pool.Exec(ctx, query, versionID, valueJSON) + _, err = db.getExecutor(tx).Exec(ctx, insertQuery, versionID, valueJSON) + if err != nil { return nil, fmt.Errorf("failed to insert server: %w", err) } @@ -367,7 +381,7 @@ func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON } // UpdateServer updates an existing server record with new server details -func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) { +func (db *PostgreSQL) UpdateServer(ctx context.Context, tx pgx.Tx, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } @@ -390,7 +404,7 @@ func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0 WHERE version_id = $2 ` - result, err := db.pool.Exec(ctx, query, valueJSON, id) + result, err := db.getExecutor(tx).Exec(ctx, query, valueJSON, id) if err != nil { return nil, fmt.Errorf("failed to update server: %w", err) } @@ -402,36 +416,29 @@ func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0 return server, nil } -// WithPublishLock executes a function with an exclusive advisory lock for publishing a server -// This prevents race conditions when multiple versions are published concurrently -func (db *PostgreSQL) WithPublishLock(ctx context.Context, serverName string, fn func(ctx context.Context) error) error { +// InTransaction executes a function within a database transaction +func (db *PostgreSQL) InTransaction(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) error { if ctx.Err() != nil { return ctx.Err() } - // Begin a transaction tx, err := db.pool.Begin(ctx) if err != nil { return fmt.Errorf("failed to begin transaction: %w", err) } + //nolint:contextcheck // Intentionally using separate context for rollback to ensure cleanup even if request is cancelled defer func() { - _ = tx.Rollback(ctx) + rollbackCtx, cancel := context.WithTimeout(context.Background(), 1*time.Second) + defer cancel() + if rbErr := tx.Rollback(rollbackCtx); rbErr != nil && !errors.Is(rbErr, pgx.ErrTxClosed) { + log.Printf("failed to rollback transaction: %v", rbErr) + } }() - // Acquire advisory lock based on server name hash - // Using pg_advisory_xact_lock which auto-releases on transaction end - lockID := hashServerName(serverName) - _, err = tx.Exec(ctx, "SELECT pg_advisory_xact_lock($1)", lockID) - if err != nil { - return fmt.Errorf("failed to acquire publish lock: %w", err) - } - - // Execute the function - if err := fn(ctx); err != nil { + if err := fn(ctx, tx); err != nil { return err } - // Commit the transaction (which also releases the lock) if err := tx.Commit(ctx); err != nil { return fmt.Errorf("failed to commit transaction: %w", err) } @@ -439,10 +446,26 @@ func (db *PostgreSQL) WithPublishLock(ctx context.Context, serverName string, fn return nil } +// AcquirePublishLock acquires an exclusive advisory lock for publishing a server +// This prevents race conditions when multiple versions are published concurrently +// Using pg_advisory_xact_lock which auto-releases on transaction end +func (db *PostgreSQL) AcquirePublishLock(ctx context.Context, tx pgx.Tx, serverName string) error { + if ctx.Err() != nil { + return ctx.Err() + } + + lockID := hashServerName(serverName) + + if _, err := db.getExecutor(tx).Exec(ctx, "SELECT pg_advisory_xact_lock($1)", lockID); err != nil { + return fmt.Errorf("failed to acquire publish lock: %w", err) + } + + return nil +} + // hashServerName creates a consistent hash of the server name for advisory locking // We use FNV-1a hash and mask to 63 bits to fit in PostgreSQL's bigint range func hashServerName(name string) int64 { - // FNV-1a 64-bit hash const ( offset64 = 14695981039346656037 prime64 = 1099511628211 @@ -452,7 +475,6 @@ func hashServerName(name string) int64 { hash ^= uint64(name[i]) hash *= prime64 } - // Use only 63 bits to ensure positive int64 //nolint:gosec // Intentional conversion with masking to 63 bits return int64(hash & 0x7FFFFFFFFFFFFFFF) } diff --git a/internal/database/testutil.go b/internal/database/testutil.go index eecbbc0..c67e3a1 100644 --- a/internal/database/testutil.go +++ b/internal/database/testutil.go @@ -2,42 +2,110 @@ package database import ( "context" + "crypto/rand" + "encoding/binary" + "errors" "fmt" "testing" "time" + "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/stretchr/testify/require" ) -const ( - // Advisory lock key for database schema initialization - // Using a fixed key ensures all test processes coordinate on the same lock - testSchemaLockKey = 123456789 -) +const templateDBName = "mcp_registry_test_template" -// NewTestDB creates a new PostgreSQL database connection for testing. -// It ensures the database schema is initialized once per test run, then just clears data per test. +// ensureTemplateDB creates a template database with migrations applied +// Multiple processes may call this, so we handle race conditions +func ensureTemplateDB(ctx context.Context, adminConn *pgx.Conn) error { + // Check if template exists + var exists bool + err := adminConn.QueryRow(ctx, "SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)", templateDBName).Scan(&exists) + if err != nil { + return fmt.Errorf("failed to check template database: %w", err) + } + + if exists { + // Template already exists + return nil + } + + // Create template database + _, err = adminConn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s", templateDBName)) + if err != nil { + // Ignore duplicate database name error - another process created it concurrently + var pgErr *pgconn.PgError + if errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "pg_database_datname_index" { + return nil + } + return fmt.Errorf("failed to create template database: %w", err) + } + + // Connect to template and run migrations + templateURI := fmt.Sprintf("postgres://mcpregistry:mcpregistry@localhost:5432/%s?sslmode=disable", templateDBName) + templateDB, err := NewPostgreSQL(ctx, templateURI) + if err != nil { + return fmt.Errorf("failed to connect to template database: %w", err) + } + defer templateDB.Close() + + // Migrations run automatically in NewPostgreSQL + return nil +} + +// NewTestDB creates an isolated PostgreSQL database for each test by copying a template. +// The template database has migrations pre-applied, so each test is fast. // Requires PostgreSQL to be running on localhost:5432 (e.g., via docker-compose). func NewTestDB(t *testing.T) Database { t.Helper() - // Create context with timeout for database operations - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) defer cancel() - // Connect to test database - connectionURI := "postgres://mcpregistry:mcpregistry@localhost:5432/mcp-registry?sslmode=disable" - db, err := NewPostgreSQL(ctx, connectionURI) - require.NoError(t, err, "Failed to connect to test PostgreSQL database. Make sure PostgreSQL is running via: docker-compose up -d postgres") + // Connect to postgres database + adminURI := "postgres://mcpregistry:mcpregistry@localhost:5432/postgres?sslmode=disable" + adminConn, err := pgx.Connect(ctx, adminURI) + require.NoError(t, err, "Failed to connect to PostgreSQL. Make sure PostgreSQL is running via: docker-compose up -d postgres") + defer adminConn.Close(ctx) - // Initialize schema once per test suite run using advisory locks for cross-process coordination - err = initializeTestSchemaWithLock(db) - require.NoError(t, err, "Failed to initialize test database schema") + // Ensure template database exists with migrations + err = ensureTemplateDB(ctx, adminConn) + require.NoError(t, err, "Failed to initialize template database") - // Clear data for this specific test - clearTestData(t, db) + // Generate unique database name for this test + var randomBytes [8]byte + _, err = rand.Read(randomBytes[:]) + require.NoError(t, err, "Failed to generate random database id") + randomInt := binary.BigEndian.Uint64(randomBytes[:]) + dbName := fmt.Sprintf("test_%d", randomInt) - // Register cleanup function to close database connection + // Create test database from template (fast - just copies files) + _, err = adminConn.Exec(ctx, fmt.Sprintf("CREATE DATABASE %s TEMPLATE %s", dbName, templateDBName)) + require.NoError(t, err, "Failed to create test database from template") + + // Register cleanup to drop database + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cleanupCancel() + + // Terminate any remaining connections + _, _ = adminConn.Exec(cleanupCtx, fmt.Sprintf( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '%s' AND pid <> pg_backend_pid()", + dbName, + )) + + // Drop database + _, _ = adminConn.Exec(cleanupCtx, fmt.Sprintf("DROP DATABASE IF EXISTS %s", dbName)) + }) + + // Connect to test database (no migrations needed - copied from template) + testURI := fmt.Sprintf("postgres://mcpregistry:mcpregistry@localhost:5432/%s?sslmode=disable", dbName) + + db, err := NewPostgreSQL(ctx, testURI) + require.NoError(t, err, "Failed to connect to test database") + + // Register cleanup to close connection t.Cleanup(func() { if err := db.Close(); err != nil { t.Logf("Warning: failed to close test database connection: %v", err) @@ -46,92 +114,3 @@ func NewTestDB(t *testing.T) Database { return db } - -// initializeTestSchemaWithLock sets up a fresh database schema with all migrations applied -// Uses PostgreSQL advisory locks to ensure only one process initializes the schema -func initializeTestSchemaWithLock(db Database) error { - // Cast to PostgreSQL to access the connection pool - pgDB, ok := db.(*PostgreSQL) - if !ok { - return fmt.Errorf("expected PostgreSQL database instance") - } - - ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) - defer cancel() - - // Acquire advisory lock to coordinate schema initialization across processes - _, err := pgDB.pool.Exec(ctx, "SELECT pg_advisory_lock($1)", testSchemaLockKey) - if err != nil { - return fmt.Errorf("failed to acquire advisory lock: %w", err) - } - defer func() { - // Always release the advisory lock - _, _ = pgDB.pool.Exec(context.Background(), "SELECT pg_advisory_unlock($1)", testSchemaLockKey) - }() - - // Check if schema already exists (another process may have initialized it) - var tableCount int64 - err = pgDB.pool.QueryRow(ctx, "SELECT count(*) FROM information_schema.tables WHERE table_schema = 'public' AND table_name = 'servers'").Scan(&tableCount) - if err != nil { - return fmt.Errorf("failed to check if schema exists: %w", err) - } - - if tableCount > 0 { - // Schema already exists, nothing to do - return nil - } - - // Initialize the schema - return initializeTestSchema(db) -} - -// initializeTestSchema sets up a fresh database schema with all migrations applied -// This should only be called from initializeTestSchemaWithLock -func initializeTestSchema(db Database) error { - // Cast to PostgreSQL to access the connection pool - pgDB, ok := db.(*PostgreSQL) - if !ok { - return fmt.Errorf("expected PostgreSQL database instance") - } - - ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) - defer cancel() - - // Drop and recreate schema completely fresh - _, err := pgDB.pool.Exec(ctx, "DROP SCHEMA public CASCADE; CREATE SCHEMA public;") - if err != nil { - return fmt.Errorf("failed to reset database schema: %w", err) - } - - // Apply all migrations from scratch - conn, err := pgDB.pool.Acquire(ctx) - if err != nil { - return fmt.Errorf("failed to acquire connection for migration: %w", err) - } - defer conn.Release() - - migrator := NewMigrator(conn.Conn()) - err = migrator.Migrate(ctx) - if err != nil { - return fmt.Errorf("failed to run database migrations: %w", err) - } - - return nil -} - -// clearTestData removes all data from test tables while preserving schema -// This runs before each individual test -func clearTestData(t *testing.T, db Database) { - t.Helper() - - // Cast to PostgreSQL to access the connection pool - pgDB, ok := db.(*PostgreSQL) - require.True(t, ok, "Expected PostgreSQL database instance") - - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) - defer cancel() - - // Clear all data but keep schema intact - _, err := pgDB.pool.Exec(ctx, "TRUNCATE TABLE servers RESTART IDENTITY CASCADE") - require.NoError(t, err, "Failed to clear test data") -} \ No newline at end of file diff --git a/internal/importer/importer.go b/internal/importer/importer.go index 30895c3..fd63db5 100644 --- a/internal/importer/importer.go +++ b/internal/importer/importer.go @@ -37,7 +37,7 @@ func (s *Service) ImportFromPath(ctx context.Context, path string) error { // Import each server using CreateServer for _, server := range servers { - _, err := s.db.CreateServer(ctx, server) + _, err := s.db.CreateServer(ctx, nil, server) if err != nil { return fmt.Errorf("failed to import server %s: %w", server.Name, err) } @@ -181,4 +181,4 @@ func convertServerResponseToRecord(response apiv0.ServerJSON) *apiv0.ServerJSON // The response is already in the correct flattened format // Just return a pointer to it return &response -} \ No newline at end of file +} diff --git a/internal/importer/importer_test.go b/internal/importer/importer_test.go index 284e89e..04b4d9e 100644 --- a/internal/importer/importer_test.go +++ b/internal/importer/importer_test.go @@ -57,7 +57,7 @@ func TestImportService_LocalFile(t *testing.T) { require.NoError(t, err) // Verify the server was imported - servers, _, err := testDB.List(context.Background(), nil, "", 10) + servers, _, err := testDB.List(context.Background(), nil, nil, "", 10) require.NoError(t, err) assert.Len(t, servers, 1) assert.Equal(t, "io.github.test/test-server-1", servers[0].Name) @@ -101,7 +101,7 @@ func TestImportService_HTTPFile(t *testing.T) { require.NoError(t, err) // Verify the server was imported - servers, _, err := testDB.List(context.Background(), nil, "", 10) + servers, _, err := testDB.List(context.Background(), nil, nil, "", 10) require.NoError(t, err) assert.Len(t, servers, 1) assert.Equal(t, "io.github.test/http-test-server", servers[0].Name) @@ -187,7 +187,7 @@ func TestImportService_RegistryAPI(t *testing.T) { require.NoError(t, err) // Verify both servers were imported - servers, _, err := testDB.List(context.Background(), nil, "", 10) + servers, _, err := testDB.List(context.Background(), nil, nil, "", 10) require.NoError(t, err) assert.Len(t, servers, 2) diff --git a/internal/service/registry_service.go b/internal/service/registry_service.go index a2f32ed..b6dbe56 100644 --- a/internal/service/registry_service.go +++ b/internal/service/registry_service.go @@ -7,6 +7,7 @@ import ( "time" "github.com/google/uuid" + "github.com/jackc/pgx/v5" "github.com/modelcontextprotocol/registry/internal/config" "github.com/modelcontextprotocol/registry/internal/database" "github.com/modelcontextprotocol/registry/internal/validators" @@ -41,7 +42,7 @@ func (s *registryServiceImpl) List(filter *database.ServerFilter, cursor string, } // Use the database's ListServers method with pagination and filtering - serverRecords, nextCursor, err := s.db.List(ctx, filter, cursor, limit) + serverRecords, nextCursor, err := s.db.List(ctx, nil, filter, cursor, limit) if err != nil { return nil, "", err } @@ -61,7 +62,7 @@ func (s *registryServiceImpl) GetByVersionID(versionID string) (*apiv0.ServerJSO ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - serverRecord, err := s.db.GetByVersionID(ctx, versionID) + serverRecord, err := s.db.GetByVersionID(ctx, nil, versionID) if err != nil { return nil, err } @@ -76,7 +77,7 @@ func (s *registryServiceImpl) GetByServerID(serverID string) (*apiv0.ServerJSON, ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - serverRecord, err := s.db.GetByServerID(ctx, serverID) + serverRecord, err := s.db.GetByServerID(ctx, nil, serverID) if err != nil { return nil, err } @@ -91,7 +92,7 @@ func (s *registryServiceImpl) GetByServerIDAndVersion(serverID string, version s ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - serverRecord, err := s.db.GetByServerIDAndVersion(ctx, serverID, version) + serverRecord, err := s.db.GetByServerIDAndVersion(ctx, nil, serverID, version) if err != nil { return nil, err } @@ -106,7 +107,7 @@ func (s *registryServiceImpl) GetAllVersionsByServerID(serverID string) ([]apiv0 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - serverRecords, err := s.db.GetAllVersionsByServerID(ctx, serverID) + serverRecords, err := s.db.GetAllVersionsByServerID(ctx, nil, serverID) if err != nil { return nil, err } @@ -131,18 +132,22 @@ func (s *registryServiceImpl) Publish(req apiv0.ServerJSON) (*apiv0.ServerJSON, return nil, err } - // Acquire advisory lock for this server name to prevent race conditions - result, err := database.WithPublishLockT(ctx, s.db, req.Name, func(lockCtx context.Context) (*apiv0.ServerJSON, error) { + result, err := database.InTransactionT(ctx, s.db, func(txCtx context.Context, tx pgx.Tx) (*apiv0.ServerJSON, error) { publishTime := time.Now() serverJSON := req + // Acquire advisory lock to prevent concurrent publishes of the same server + if err := s.db.AcquirePublishLock(txCtx, tx, serverJSON.Name); err != nil { + return nil, err + } + // Check for duplicate remote URLs - if err := s.validateNoDuplicateRemoteURLs(lockCtx, serverJSON); err != nil { + if err := s.validateNoDuplicateRemoteURLs(txCtx, tx, serverJSON); err != nil { return nil, err } filter := &database.ServerFilter{Name: &serverJSON.Name} - existingServerVersions, _, err := s.db.List(lockCtx, filter, "", maxServerVersionsPerServer) + existingServerVersions, _, err := s.db.List(txCtx, tx, filter, "", maxServerVersionsPerServer) if err != nil && !errors.Is(err, database.ErrNotFound) { return nil, err } @@ -176,18 +181,11 @@ func (s *registryServiceImpl) Publish(req apiv0.ServerJSON) (*apiv0.ServerJSON, ) > 0 } - // Mark previous latest as no longer latest BEFORE creating new version - // This prevents violating the unique constraint on isLatest + // Unmark old latest version if needed if isNewLatest && existingLatest != nil { - var existingLatestVersionID string if existingLatest.Meta != nil && existingLatest.Meta.Official != nil { - existingLatestVersionID = existingLatest.Meta.Official.VersionID - } - if existingLatestVersionID != "" { - // Update the existing server to set isLatest = false existingLatest.Meta.Official.IsLatest = false - existingLatest.Meta.Official.UpdatedAt = time.Now() - if _, err := s.db.UpdateServer(lockCtx, existingLatestVersionID, existingLatest); err != nil { + if _, err := s.db.UpdateServer(txCtx, tx, existingLatest.Meta.Official.VersionID, existingLatest); err != nil { return nil, err } } @@ -196,13 +194,8 @@ func (s *registryServiceImpl) Publish(req apiv0.ServerJSON) (*apiv0.ServerJSON, // Create complete server with metadata server := s.createServerWithMetadata(serverJSON, existingServerVersions, publishTime, isNewLatest) - // Create server in database - serverRecord, err := s.db.CreateServer(lockCtx, &server) - if err != nil { - return nil, err - } - - return serverRecord, nil + // Insert new server version + return s.db.CreateServer(txCtx, tx, &server) }) if err != nil { @@ -253,13 +246,13 @@ func (s *registryServiceImpl) createServerWithMetadata( } // validateNoDuplicateRemoteURLs checks that no other server is using the same remote URLs -func (s *registryServiceImpl) validateNoDuplicateRemoteURLs(ctx context.Context, serverDetail apiv0.ServerJSON) error { +func (s *registryServiceImpl) validateNoDuplicateRemoteURLs(ctx context.Context, tx pgx.Tx, serverDetail apiv0.ServerJSON) error { // Check each remote URL in the new server for conflicts for _, remote := range serverDetail.Remotes { // Use filter to find servers with this remote URL filter := &database.ServerFilter{RemoteURL: &remote.URL} - conflictingServers, _, err := s.db.List(ctx, filter, "", 1000) + conflictingServers, _, err := s.db.List(ctx, tx, filter, "", 1000) if err != nil { return fmt.Errorf("failed to check remote URL conflict: %w", err) } @@ -291,45 +284,46 @@ func (s *registryServiceImpl) EditServer(versionID string, req apiv0.ServerJSON) ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() - // First get the current server to preserve metadata - currentServer, err := s.db.GetByVersionID(ctx, versionID) - if err != nil { - return nil, err - } - // Validate the request if err := validators.ValidatePublishRequest(req, s.cfg); err != nil { return nil, err } - // Merge the request with the current server, preserving metadata - updatedServer := *currentServer // Copy the current server with all metadata + return database.InTransactionT(ctx, s.db, func(txCtx context.Context, tx pgx.Tx) (*apiv0.ServerJSON, error) { + // Get the current server to preserve metadata + currentServer, err := s.db.GetByVersionID(txCtx, tx, versionID) + if err != nil { + return nil, err + } - // Update only the user-modifiable fields from the request - updatedServer.Name = req.Name - updatedServer.Description = req.Description - updatedServer.Version = req.Version - updatedServer.Status = req.Status - updatedServer.Repository = req.Repository - updatedServer.Remotes = req.Remotes - updatedServer.Packages = req.Packages + // Acquire advisory lock to prevent concurrent edits of servers with same name + if err := s.db.AcquirePublishLock(txCtx, tx, currentServer.Name); err != nil { + return nil, err + } - // Update the UpdatedAt timestamp in metadata - if updatedServer.Meta != nil && updatedServer.Meta.Official != nil { - updatedServer.Meta.Official.UpdatedAt = time.Now() - } + // Merge the request with the current server, preserving metadata + updatedServer := *currentServer - // Check for duplicate remote URLs using the updated server - if err := s.validateNoDuplicateRemoteURLs(ctx, updatedServer); err != nil { - return nil, err - } + // Update only the user-modifiable fields from the request + updatedServer.Name = req.Name + updatedServer.Description = req.Description + updatedServer.Version = req.Version + updatedServer.Status = req.Status + updatedServer.Repository = req.Repository + updatedServer.Remotes = req.Remotes + updatedServer.Packages = req.Packages - // Update server in database - serverRecord, err := s.db.UpdateServer(ctx, versionID, &updatedServer) - if err != nil { - return nil, err - } + // Update the UpdatedAt timestamp in metadata + if updatedServer.Meta != nil && updatedServer.Meta.Official != nil { + updatedServer.Meta.Official.UpdatedAt = time.Now() + } - // Return the server record directly - return serverRecord, nil + // Check for duplicate remote URLs using the updated server + if err := s.validateNoDuplicateRemoteURLs(txCtx, tx, updatedServer); err != nil { + return nil, err + } + + // Update server in database + return s.db.UpdateServer(txCtx, tx, versionID, &updatedServer) + }) } diff --git a/internal/service/registry_service_test.go b/internal/service/registry_service_test.go index 48f2f4b..29743f8 100644 --- a/internal/service/registry_service_test.go +++ b/internal/service/registry_service_test.go @@ -107,7 +107,7 @@ func TestValidateNoDuplicateRemoteURLs(t *testing.T) { ctx := context.Background() impl := service.(*registryServiceImpl) - err := impl.validateNoDuplicateRemoteURLs(ctx, tt.serverDetail) + err := impl.validateNoDuplicateRemoteURLs(ctx, nil, tt.serverDetail) if tt.expectError { assert.Error(t, err) @@ -384,7 +384,7 @@ func TestPublishConcurrentVersionsNoRace(t *testing.T) { testDB := database.NewTestDB(t) service := NewRegistryService(testDB, &config.Config{EnableRegistryValidation: false}) - const concurrency = 1 // @Maintainers: Fix this and increase to higher number, previously 100 + const concurrency = 100 results := make([]*apiv0.ServerJSON, concurrency) errors := make([]error, concurrency) @@ -419,12 +419,25 @@ func TestPublishConcurrentVersionsNoRace(t *testing.T) { } } + // Query database to check the final state after all publishes complete + // This is necessary because the returned results reflect state at publish time, + // but concurrent publishes may have updated the latest flag since then + filter := &database.ServerFilter{Name: strPtr("com.example/test-concurrent")} + dbResults, _, err := testDB.List(context.Background(), nil, filter, "", 1000) + assert.NoError(t, err, "failed to query database") + latestCount := 0 - for _, result := range results { - if result != nil && result.Meta != nil && result.Meta.Official != nil && - result.Meta.Official.IsLatest { + var latestVersion string + for _, r := range dbResults { + if r.Meta != nil && r.Meta.Official != nil && r.Meta.Official.IsLatest { latestCount++ + latestVersion = r.Version } } - assert.Equal(t, 1, latestCount, "should have exactly one latest version") + + assert.Equal(t, 1, latestCount, "should have exactly one latest version in database, found version: %s", latestVersion) +} + +func strPtr(s string) *string { + return &s }