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
1 changed files with 846 additions and 564 deletions
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