Use a connection pool instead of a single DB connection (#408)
<!-- Provide a brief summary of your changes --> ## Motivation and Context <!-- Why is this change needed? What problem does it solve? --> The following PR should fix the recent "conn busy" issue caused by the PostgreSQL implementation using a single database connection (*pgx.Conn) instead of a connection pool (*pgxpool.Pool). When multiple concurrent requests try to access the endpoint, they would conflict trying to use the same connection. ## How Has This Been Tested? <!-- Have you tested this in a real application? Which scenarios were tested? --> Builds locally ## Breaking Changes <!-- Will users need to update their code or configurations? --> No ## Types of changes <!-- What types of changes does your code introduce? Put an `x` in all the boxes that apply: --> - [x] Bug fix (non-breaking change which fixes an issue) - [ ] New feature (non-breaking change which adds functionality) - [ ] 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. --> - [ ] I have read the [MCP Documentation](https://modelcontextprotocol.io) - [ ] My code follows the repository's style guidelines - [ ] New and existing tests pass locally - [ ] I have added appropriate error handling - [ ] I have added or updated documentation as needed ## Additional context <!-- Add any other context, implementation notes, or design decisions --> Related to #400, #405 and #407 --------- Signed-off-by: Radoslav Dimitrov <radoslav@stacklok.com>
This commit is contained in:
parent
ca3da9045b
commit
e6b00ee0ed
1 changed files with 36 additions and 14 deletions
|
|
@ -6,37 +6,58 @@ import (
|
|||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
|
||||
)
|
||||
|
||||
// PostgreSQL is an implementation of the Database interface using PostgreSQL
|
||||
type PostgreSQL struct {
|
||||
conn *pgx.Conn
|
||||
pool *pgxpool.Pool
|
||||
}
|
||||
|
||||
// NewPostgreSQL creates a new instance of the PostgreSQL database
|
||||
func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, error) {
|
||||
conn, err := pgx.Connect(ctx, connectionURI)
|
||||
// Parse connection config for pool settings
|
||||
config, err := pgxpool.ParseConfig(connectionURI)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to connect to PostgreSQL: %w", err)
|
||||
return nil, fmt.Errorf("failed to parse PostgreSQL config: %w", err)
|
||||
}
|
||||
|
||||
// Configure pool for stability-focused defaults
|
||||
config.MaxConns = 30 // Handle good concurrent load
|
||||
config.MinConns = 5 // Keep connections warm for fast response
|
||||
config.MaxConnIdleTime = 30 * time.Minute // Keep connections available for bursts
|
||||
config.MaxConnLifetime = 2 * time.Hour // Refresh connections regularly for stability
|
||||
|
||||
// Create connection pool with configured settings
|
||||
pool, err := pgxpool.NewWithConfig(ctx, config)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create PostgreSQL pool: %w", err)
|
||||
}
|
||||
|
||||
// Test the connection
|
||||
if err = conn.Ping(ctx); err != nil {
|
||||
if err = pool.Ping(ctx); err != nil {
|
||||
return nil, fmt.Errorf("failed to ping PostgreSQL: %w", err)
|
||||
}
|
||||
|
||||
// Run migrations
|
||||
migrator := NewMigrator(conn)
|
||||
// Run migrations using a single connection from the pool
|
||||
conn, err := pool.Acquire(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to acquire connection for migrations: %w", err)
|
||||
}
|
||||
defer conn.Release()
|
||||
|
||||
migrator := NewMigrator(conn.Conn())
|
||||
if err := migrator.Migrate(ctx); err != nil {
|
||||
return nil, fmt.Errorf("failed to run database migrations: %w", err)
|
||||
}
|
||||
|
||||
return &PostgreSQL{
|
||||
conn: conn,
|
||||
pool: pool,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
|
@ -120,7 +141,7 @@ func (db *PostgreSQL) List(
|
|||
`, whereClause, argIndex)
|
||||
args = append(args, limit)
|
||||
|
||||
rows, err := db.conn.Query(ctx, query, args...)
|
||||
rows, err := db.pool.Query(ctx, query, args...)
|
||||
if err != nil {
|
||||
return nil, "", fmt.Errorf("failed to query servers: %w", err)
|
||||
}
|
||||
|
|
@ -173,7 +194,7 @@ func (db *PostgreSQL) GetByID(ctx context.Context, id string) (*apiv0.ServerJSON
|
|||
|
||||
var valueJSON []byte
|
||||
|
||||
err := db.conn.QueryRow(ctx, query, id).Scan(&valueJSON)
|
||||
err := db.pool.QueryRow(ctx, query, id).Scan(&valueJSON)
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, pgx.ErrNoRows) {
|
||||
|
|
@ -216,7 +237,7 @@ func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON
|
|||
VALUES ($1, $2)
|
||||
`
|
||||
|
||||
_, err = db.conn.Exec(ctx, query, id, valueJSON)
|
||||
_, err = db.pool.Exec(ctx, query, id, valueJSON)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to insert server: %w", err)
|
||||
}
|
||||
|
|
@ -248,7 +269,7 @@ func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0
|
|||
WHERE id = $2
|
||||
`
|
||||
|
||||
result, err := db.conn.Exec(ctx, query, valueJSON, id)
|
||||
result, err := db.pool.Exec(ctx, query, valueJSON, id)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to update server: %w", err)
|
||||
}
|
||||
|
|
@ -262,5 +283,6 @@ func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0
|
|||
|
||||
// Close closes the database connection
|
||||
func (db *PostgreSQL) Close() error {
|
||||
return db.conn.Close(context.Background())
|
||||
db.pool.Close()
|
||||
return nil
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue