Fix atomic latest version update to prevent missing isLatest flags (#530)
## 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 <noreply@anthropic.com>
This commit is contained in:
adam jones
GitHub
parent
ae00f70eaf
commit
747ad310bb
1 changed files with 243 additions and 233 deletions
|
|
@ -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")
|
||||
}
|
||||
Loading…
Reference in a new issue