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:
Radoslav Dimitrov 2025-09-10 16:45:30 +03:00 • committed by GitHub
commit e6b00ee0ed
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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
}