Database implementation cleanup (#335)

## Motivation and Context

The database logic was getting a bit unwieldy and fragile - especially
because business logic that should really have lived at the registry
service level was being pushed into the persistence level.

This PR tightens the interface, simplifies the code there, as well as
rips out a bunch of mocks etc. we don't need.

## How Has This Been Tested?

Locally everything works

## Breaking Changes

The list endpoint was surprisingly always filtering for isLatest only,
which seems like a bug? I have removed this filter.
This commit is contained in:
adam jones 2025-09-02 13:37:28 +01:00 • committed by GitHub
commit 9d9a02c57f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 669 additions and 1280 deletions

View file

@ -17,21 +17,22 @@ var (
ErrMaxServersReached = errors.New("maximum number of versions for this server reached (10000): please reach out at https://github.com/modelcontextprotocol/registry to explain your use case")
)
// Database defines the interface for database operations with extension wrapper architecture
// ServerFilter defines filtering options for server queries
type ServerFilter struct {
Name *string // for finding versions of same server
RemoteURL *string // for duplicate URL detection
}
// Database defines the interface for database operations
type Database interface {
// List retrieves all ServerRecord entries with optional filtering
List(ctx context.Context, filter map[string]any, cursor string, limit int) ([]*apiv0.ServerJSON, string, error)
// GetByID retrieves a single ServerRecord by its ID
// Retrieve server entries with optional filtering
List(ctx context.Context, filter *ServerFilter, cursor string, limit int) ([]*apiv0.ServerJSON, string, error)
// Retrieve a single server by its ID
GetByID(ctx context.Context, id string) (*apiv0.ServerJSON, error)
// Publish adds a new server to the database with separated server.json and extensions
// The registryMetadata contains metadata determined by the service layer (e.g., is_latest, timestamps)
Publish(ctx context.Context, serverDetail apiv0.ServerJSON, publisherExtensions map[string]interface{}, registryMetadata apiv0.RegistryExtensions) (*apiv0.ServerJSON, error)
// UpdateLatestFlag updates the is_latest flag for a specific server record
UpdateLatestFlag(ctx context.Context, id string, isLatest bool) error
// UpdateServer updates an existing server record with new server details
UpdateServer(ctx context.Context, id string, serverDetail apiv0.ServerJSON, publisherExtensions map[string]interface{}) (*apiv0.ServerJSON, error)
// ImportSeed imports initial data from a seed file
ImportSeed(ctx context.Context, seedFilePath string) error
// CreateServer adds a new server to the database
CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error)
// UpdateServer updates an existing server record
UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error)
// Close closes the database connection
Close() error
}
@ -45,15 +46,3 @@ const (
// ConnectionTypePostgreSQL represents a PostgreSQL database connection
ConnectionTypePostgreSQL ConnectionType = "postgresql"
)
// ConnectionInfo provides information about the database connection
type ConnectionInfo struct {
// Type indicates the type of database connection
Type ConnectionType
// IsConnected indicates whether the database is currently connected
IsConnected bool
// Raw provides access to the underlying connection object, which will vary by implementation
// For PostgreSQL, this will be *pgx.Conn
// For MemoryDB, this will be map[string]*model.MCPRegistry
Raw any
}

View file

@ -1,155 +0,0 @@
package database
import (
"context"
"encoding/json"
"fmt"
"io"
"log"
"net/http"
"os"
"strings"
"github.com/modelcontextprotocol/registry/internal/validators"
apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
)
// ReadSeedFile reads seed data from various sources:
// 1. Local file paths (*.json files) - expects extension wrapper format
// 2. Direct HTTP URLs to seed.json files - expects extension wrapper format
// 3. Registry root URLs (automatically appends /v0/servers and paginates)
// Only the extension wrapper format is supported (array of ServerResponse objects)
func ReadSeedFile(ctx context.Context, path string) ([]*apiv0.ServerJSON, error) {
var data []byte
var err error
if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
// Handle HTTP URLs
if strings.HasSuffix(path, "/v0/servers") || strings.Contains(path, "/v0/servers") {
// This is a registry API endpoint - fetch paginated data
return fetchFromRegistryAPI(ctx, path)
}
// This is a direct file URL
data, err = fetchFromHTTP(ctx, path)
} else {
// Handle local file paths
data, err = os.ReadFile(path)
}
if err != nil {
return nil, fmt.Errorf("failed to read seed data from %s: %w", path, err)
}
// Parse extension wrapper format (only supported format)
var serverResponses []apiv0.ServerJSON
if err := json.Unmarshal(data, &serverResponses); err != nil {
return nil, fmt.Errorf("failed to parse seed data as extension wrapper format: %w", err)
}
if len(serverResponses) == 0 {
return []*apiv0.ServerJSON{}, nil
}
// Validate servers and collect warnings instead of failing the whole batch
var validRecords []*apiv0.ServerJSON
var invalidServers []string
var validationFailures []string
for _, response := range serverResponses {
if err := validators.ValidateServerJSON(&response); err != nil {
// Log warning and track invalid server instead of failing
invalidServers = append(invalidServers, response.Name)
validationFailures = append(validationFailures, fmt.Sprintf("Server '%s': %v", response.Name, err))
log.Printf("Warning: Skipping invalid server '%s': %v", response.Name, err)
continue
}
// Convert valid ServerJSON to ServerRecord
record := convertServerResponseToRecord(response)
validRecords = append(validRecords, record)
}
// Print summary of validation results
if len(invalidServers) > 0 {
log.Printf("Import summary: %d valid servers imported, %d invalid servers skipped", len(validRecords), len(invalidServers))
log.Printf("Invalid servers: %v", invalidServers)
for _, failure := range validationFailures {
log.Printf(" - %s", failure)
}
} else {
log.Printf("Import summary: All %d servers imported successfully", len(validRecords))
}
return validRecords, nil
}
func fetchFromHTTP(ctx context.Context, url string) ([]byte, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP request: %w", err)
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to fetch from HTTP: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("HTTP request failed with status: %d", resp.StatusCode)
}
return io.ReadAll(resp.Body)
}
func fetchFromRegistryAPI(ctx context.Context, baseURL string) ([]*apiv0.ServerJSON, error) {
var allRecords []*apiv0.ServerJSON
cursor := ""
for {
url := baseURL
if cursor != "" {
if strings.Contains(url, "?") {
url += "&cursor=" + cursor
} else {
url += "?cursor=" + cursor
}
}
data, err := fetchFromHTTP(ctx, url)
if err != nil {
return nil, fmt.Errorf("failed to fetch page from registry API: %w", err)
}
var response struct {
Servers []apiv0.ServerJSON `json:"servers"`
Metadata *struct {
NextCursor string `json:"next_cursor,omitempty"`
} `json:"metadata,omitempty"`
}
if err := json.Unmarshal(data, &response); err != nil {
return nil, fmt.Errorf("failed to parse registry API response: %w", err)
}
// Convert and add servers
for _, serverResponse := range response.Servers {
record := convertServerResponseToRecord(serverResponse)
allRecords = append(allRecords, record)
}
// Check if there's a next page
if response.Metadata == nil || response.Metadata.NextCursor == "" {
break
}
cursor = response.Metadata.NextCursor
}
return allRecords, nil
}
func convertServerResponseToRecord(response apiv0.ServerJSON) *apiv0.ServerJSON {
// The response is already in the correct flattened format
// Just return a pointer to it
return &response
}

View file

@ -1,268 +0,0 @@
package database_test
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"testing"
"time"
"github.com/modelcontextprotocol/registry/internal/database"
apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
"github.com/modelcontextprotocol/registry/pkg/model"
"github.com/stretchr/testify/assert"
)
func TestReadSeedFile_LocalFile(t *testing.T) {
// Create a temporary seed file in extension wrapper format
tempFile := "/tmp/test_seed.json"
seedData := []apiv0.ServerJSON{
{
Name: "com.example/test-server-1",
Description: "Test server 1",
Repository: model.Repository{
URL: "https://github.com/test/repo1",
Source: "github",
ID: "123",
},
VersionDetail: model.VersionDetail{
Version: "1.0.0",
},
Remotes: []model.Remote{
{
URL: "https://example.com/remote",
},
},
Meta: &apiv0.ServerMeta{
IOModelContextProtocolRegistry: &apiv0.RegistryExtensions{
ID: "test-id-1",
PublishedAt: time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC),
IsLatest: true,
ReleaseDate: "2023-01-01T00:00:00Z",
},
},
},
}
// Write seed data to temp file
data, err := json.Marshal(seedData)
assert.NoError(t, err)
err = func() error {
f, err := os.Create(tempFile)
if err != nil {
return err
}
defer f.Close()
_, err = f.Write(data)
return err
}()
assert.NoError(t, err)
defer os.Remove(tempFile)
// Test reading the file
result, err := database.ReadSeedFile(context.Background(), tempFile)
assert.NoError(t, err)
assert.Len(t, result, 1)
assert.Equal(t, "com.example/test-server-1", result[0].Name)
}
func TestReadSeedFile_DirectHTTPURL(t *testing.T) {
// Create a test HTTP server that serves seed JSON directly in extension wrapper format
seedData := []apiv0.ServerJSON{
{
Name: "com.example/test-server-1",
Description: "Test server 1",
Remotes: []model.Remote{
{
URL: "https://example.com/remote",
},
},
Meta: &apiv0.ServerMeta{
IOModelContextProtocolRegistry: &apiv0.RegistryExtensions{
ID: "test-id-1",
PublishedAt: time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC),
IsLatest: true,
ReleaseDate: "2023-01-01T00:00:00Z",
},
},
},
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(seedData); err != nil {
http.Error(w, "Failed to encode JSON", http.StatusInternalServerError)
}
}))
defer server.Close()
// Test reading from HTTP URL ending in .json
result, err := database.ReadSeedFile(context.Background(), server.URL+"/seed.json")
assert.NoError(t, err)
assert.Len(t, result, 1)
assert.Equal(t, "com.example/test-server-1", result[0].Name)
}
func TestReadSeedFile_RegistryURL(t *testing.T) {
// Create mock registry responses
server1 := apiv0.ServerJSON{
Name: "Test Server 1",
Description: "First test server",
Packages: []model.Package{
{
RegistryType: "npm",
RegistryBaseURL: "https://registry.npmjs.org",
Identifier: "test-package-1",
Version: "1.0.0",
},
},
Meta: &apiv0.ServerMeta{
IOModelContextProtocolRegistry: &apiv0.RegistryExtensions{
ID: "server-1",
PublishedAt: time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC),
IsLatest: true,
ReleaseDate: "2023-01-01T00:00:00Z",
},
},
}
server2 := apiv0.ServerJSON{
Name: "Test Server 2",
Description: "Second test server",
Packages: []model.Package{
{
RegistryType: "npm",
RegistryBaseURL: "https://registry.npmjs.org",
Identifier: "test-package-2",
Version: "2.0.0",
},
},
Meta: &apiv0.ServerMeta{
IOModelContextProtocolRegistry: &apiv0.RegistryExtensions{
ID: "server-2",
PublishedAt: time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC),
IsLatest: true,
ReleaseDate: "2023-01-01T00:00:00Z",
},
},
}
serverDetail1 := apiv0.ServerJSON{
Name: "Test Server 1",
Description: "First test server",
Packages: []model.Package{
{
RegistryType: "npm",
RegistryBaseURL: "https://registry.npmjs.org",
Identifier: "test-package-1",
Version: "1.0.0",
},
},
}
serverDetail2 := apiv0.ServerJSON{
Name: "Test Server 2",
Description: "Second test server",
Packages: []model.Package{
{
RegistryType: "npm",
RegistryBaseURL: "https://registry.npmjs.org",
Identifier: "test-package-2",
Version: "2.0.0",
},
},
}
// Create a test HTTP server that simulates the registry API
mux := http.NewServeMux()
// Handle /v0/servers endpoint (paginated)
mux.HandleFunc("/v0/servers", func(w http.ResponseWriter, r *http.Request) {
cursor := r.URL.Query().Get("cursor")
type Metadata struct {
NextCursor string `json:"next_cursor,omitempty"`
Count int `json:"count,omitempty"`
}
type PaginatedResponse struct {
Servers []apiv0.ServerJSON `json:"servers"`
Metadata *Metadata `json:"metadata,omitempty"`
}
var response PaginatedResponse
switch cursor {
case "":
// First page
response = PaginatedResponse{
Servers: []apiv0.ServerJSON{server1},
Metadata: &Metadata{
NextCursor: "next-cursor-1",
Count: 1,
},
}
case "next-cursor-1":
// Second page
response = PaginatedResponse{
Servers: []apiv0.ServerJSON{server2},
Metadata: &Metadata{
Count: 1,
// No NextCursor means end of pagination
},
}
default:
// No more pages
response = PaginatedResponse{
Servers: []apiv0.ServerJSON{},
Metadata: &Metadata{},
}
}
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(response); err != nil {
http.Error(w, "Failed to encode JSON", http.StatusInternalServerError)
}
})
// Handle individual server detail endpoints
mux.HandleFunc("/v0/servers/server-1", func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(serverDetail1); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
})
mux.HandleFunc("/v0/servers/server-2", func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
if err := json.NewEncoder(w).Encode(serverDetail2); err != nil {
http.Error(w, err.Error(), http.StatusInternalServerError)
}
})
server := httptest.NewServer(mux)
defer server.Close()
// Test reading from registry root URL (this should trigger pagination)
result, err := database.ReadSeedFile(context.Background(), server.URL+"/v0/servers")
assert.NoError(t, err)
assert.Len(t, result, 2)
// Verify the servers were imported correctly
assert.Equal(t, "Test Server 1", result[0].Name)
assert.Equal(t, "Test Server 2", result[1].Name)
// Verify packages were imported with new schema
assert.Len(t, result[0].Packages, 1)
assert.Equal(t, "npm", result[0].Packages[0].RegistryType)
assert.Equal(t, "https://registry.npmjs.org", result[0].Packages[0].RegistryBaseURL)
assert.Equal(t, "test-package-1", result[0].Packages[0].Identifier)
assert.Len(t, result[1].Packages, 1)
assert.Equal(t, "npm", result[1].Packages[0].RegistryType)
assert.Equal(t, "https://registry.npmjs.org", result[1].Packages[0].RegistryBaseURL)
assert.Equal(t, "test-package-2", result[1].Packages[0].Identifier)
// Verify metadata was extracted
assert.Equal(t, "server-1", result[0].Meta.IOModelContextProtocolRegistry.ID)
assert.Equal(t, "server-2", result[1].Meta.IOModelContextProtocolRegistry.ID)
}

View file

@ -5,7 +5,6 @@ import (
"fmt"
"sort"
"sync"
"time"
apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
)
@ -16,38 +15,17 @@ type MemoryDB struct {
mu sync.RWMutex
}
// NewMemoryDB creates a new instance of the in-memory database
func NewMemoryDB(e map[string]*apiv0.ServerJSON) *MemoryDB {
func NewMemoryDB() *MemoryDB {
// Convert input ServerJSON entries to have proper metadata
serverRecords := make(map[string]*apiv0.ServerJSON)
for registryID, serverDetail := range e {
// Create a copy and add registry metadata
now := time.Now()
server := *serverDetail // Copy the server
if server.Meta == nil {
server.Meta = &apiv0.ServerMeta{}
}
server.Meta.IOModelContextProtocolRegistry = &apiv0.RegistryExtensions{
ID: registryID,
PublishedAt: now,
UpdatedAt: now,
IsLatest: true,
ReleaseDate: now.Format(time.RFC3339),
}
if server.Meta.Publisher == nil {
server.Meta.Publisher = make(map[string]interface{})
}
serverRecords[registryID] = &server
}
return &MemoryDB{
entries: serverRecords,
}
}
//nolint:cyclop,gocyclo,gocognit // Complexity from filtering logic is acceptable for memory implementation
func (db *MemoryDB) List(
ctx context.Context,
filter map[string]any,
filter *ServerFilter,
cursor string,
limit int,
) ([]*apiv0.ServerJSON, string, error) {
@ -62,73 +40,20 @@ func (db *MemoryDB) List(
db.mu.RLock()
defer db.mu.RUnlock()
// Convert all entries to a slice for pagination, filter by is_latest
// Convert all entries to a slice for pagination
var allEntries []*apiv0.ServerJSON
for _, entry := range db.entries {
if entry.Meta != nil && entry.Meta.IOModelContextProtocolRegistry != nil && entry.Meta.IOModelContextProtocolRegistry.IsLatest {
allEntries = append(allEntries, entry)
}
allEntries = append(allEntries, entry)
}
// Simple filtering implementation
var filteredEntries []*apiv0.ServerJSON
for _, entry := range allEntries {
include := true
// Apply filtering and sorting
filteredEntries := db.filterAndSort(allEntries, filter)
// Apply filters if any
for key, value := range filter {
switch key {
case "name":
if entry.Name != value.(string) {
include = false
}
case "version":
if entry.VersionDetail.Version != value.(string) {
include = false
}
case "status":
if string(entry.Status) != value.(string) {
include = false
}
case "remote_url":
found := false
remoteURL := value.(string)
for _, remote := range entry.Remotes {
if remote.URL == remoteURL {
found = true
break
}
}
if !found {
include = false
}
}
}
if include {
filteredEntries = append(filteredEntries, entry)
}
}
// Sort filteredEntries by registry metadata ID for consistent pagination
sort.Slice(filteredEntries, func(i, j int) bool {
iID := ""
jID := ""
if filteredEntries[i].Meta != nil && filteredEntries[i].Meta.IOModelContextProtocolRegistry != nil {
iID = filteredEntries[i].Meta.IOModelContextProtocolRegistry.ID
}
if filteredEntries[j].Meta != nil && filteredEntries[j].Meta.IOModelContextProtocolRegistry != nil {
jID = filteredEntries[j].Meta.IOModelContextProtocolRegistry.ID
}
return iID < jID
})
// Find starting point for cursor-based pagination using registry metadata ID
// Find starting point for cursor-based pagination
startIdx := 0
if cursor != "" {
for i, entry := range filteredEntries {
if entry.Meta != nil && entry.Meta.IOModelContextProtocolRegistry != nil &&
entry.Meta.IOModelContextProtocolRegistry.ID == cursor {
if db.getRegistryID(entry) == cursor {
startIdx = i + 1 // Start after the cursor
break
}
@ -145,19 +70,15 @@ func (db *MemoryDB) List(
result = []*apiv0.ServerJSON{}
}
// Determine next cursor using registry metadata ID
// Determine next cursor
nextCursor := ""
if endIdx < len(filteredEntries) && len(result) > 0 {
lastEntry := result[len(result)-1]
if lastEntry.Meta != nil && lastEntry.Meta.IOModelContextProtocolRegistry != nil {
nextCursor = lastEntry.Meta.IOModelContextProtocolRegistry.ID
}
nextCursor = db.getRegistryID(result[len(result)-1])
}
return result, nextCursor, nil
}
// GetByID retrieves a single ServerRecord by its registry metadata ID
func (db *MemoryDB) GetByID(ctx context.Context, id string) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
@ -176,103 +97,28 @@ func (db *MemoryDB) GetByID(ctx context.Context, id string) (*apiv0.ServerJSON,
return nil, ErrNotFound
}
// Publish adds a new server to the database with separated server.json and extensions
func (db *MemoryDB) Publish(ctx context.Context, serverDetail apiv0.ServerJSON, publisherExtensions map[string]interface{}, registryMetadata apiv0.RegistryExtensions) (*apiv0.ServerJSON, error) {
func (db *MemoryDB) CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Extract name and version for validation
name := serverDetail.Name
if name == "" {
return nil, fmt.Errorf("name is required in server JSON")
// Get the ID from the registry metadata
if server.Meta == nil || server.Meta.IOModelContextProtocolRegistry == nil {
return nil, fmt.Errorf("server must have registry metadata with ID")
}
version := serverDetail.VersionDetail.Version
if version == "" {
return nil, fmt.Errorf("version is required in version_detail")
}
id := server.Meta.IOModelContextProtocolRegistry.ID
db.mu.Lock()
defer db.mu.Unlock()
// Validate repository URL
if serverDetail.Repository.URL == "" {
return nil, ErrInvalidInput
}
// Create server record with metadata
record := serverDetail // Copy the input
// Initialize meta if not present
if record.Meta == nil {
record.Meta = &apiv0.ServerMeta{}
}
// Set registry metadata
record.Meta.IOModelContextProtocolRegistry = &registryMetadata
// Set publisher extensions if provided
if len(publisherExtensions) > 0 {
record.Meta.Publisher = publisherExtensions
}
// Store the record using registry metadata ID
db.entries[registryMetadata.ID] = &record
db.entries[id] = server
return &record, nil
return server, nil
}
// ImportSeed imports initial data from a seed file into memory database
func (db *MemoryDB) ImportSeed(ctx context.Context, seedFilePath string) error {
if ctx.Err() != nil {
return ctx.Err()
}
// This will need to be updated to work with the new ServerRecord format
// Read seed data using the shared ReadSeedFile function
seedData, err := ReadSeedFile(ctx, seedFilePath)
if err != nil {
return fmt.Errorf("failed to read seed file: %w", err)
}
// Lock for concurrent access
db.mu.Lock()
defer db.mu.Unlock()
// Import each server
for _, record := range seedData {
// Use the registry metadata ID as the map key
if record.Meta != nil && record.Meta.IOModelContextProtocolRegistry != nil {
db.entries[record.Meta.IOModelContextProtocolRegistry.ID] = record
}
}
return nil
}
// UpdateLatestFlag updates the is_latest flag for a specific server record
func (db *MemoryDB) UpdateLatestFlag(ctx context.Context, id string, isLatest bool) error {
if ctx.Err() != nil {
return ctx.Err()
}
db.mu.Lock()
defer db.mu.Unlock()
if entry, exists := db.entries[id]; exists {
if entry.Meta != nil && entry.Meta.IOModelContextProtocolRegistry != nil {
entry.Meta.IOModelContextProtocolRegistry.IsLatest = isLatest
entry.Meta.IOModelContextProtocolRegistry.UpdatedAt = time.Now()
}
return nil
}
return ErrNotFound
}
// UpdateServer updates an existing server record with new server details
func (db *MemoryDB) UpdateServer(ctx context.Context, id string, serverDetail apiv0.ServerJSON, publisherExtensions map[string]interface{}) (*apiv0.ServerJSON, error) {
func (db *MemoryDB) UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
@ -280,40 +126,75 @@ func (db *MemoryDB) UpdateServer(ctx context.Context, id string, serverDetail ap
db.mu.Lock()
defer db.mu.Unlock()
entry, exists := db.entries[id]
_, exists := db.entries[id]
if !exists {
return nil, ErrNotFound
}
// Update the server details by copying fields from serverDetail
*entry = serverDetail
// Ensure meta exists and update it
if entry.Meta == nil {
entry.Meta = &apiv0.ServerMeta{}
}
if entry.Meta.IOModelContextProtocolRegistry != nil {
entry.Meta.IOModelContextProtocolRegistry.UpdatedAt = time.Now()
}
if len(publisherExtensions) > 0 {
entry.Meta.Publisher = publisherExtensions
}
// Update the server
db.entries[id] = server
// Return the updated record
return entry, nil
return server, nil
}
// Close closes the database connection
// For an in-memory database, this is a no-op
func (db *MemoryDB) Close() error {
return nil
}
// Connection returns information about the database connection
func (db *MemoryDB) Connection() *ConnectionInfo {
return &ConnectionInfo{
Type: ConnectionTypeMemory,
IsConnected: true, // Memory DB is always connected
Raw: db.entries,
// filterAndSort applies filtering and sorting to the entries
func (db *MemoryDB) filterAndSort(allEntries []*apiv0.ServerJSON, filter *ServerFilter) []*apiv0.ServerJSON {
// Apply filtering
var filteredEntries []*apiv0.ServerJSON
for _, entry := range allEntries {
if db.matchesFilter(entry, filter) {
filteredEntries = append(filteredEntries, entry)
}
}
// Sort by registry metadata ID for consistent pagination
sort.Slice(filteredEntries, func(i, j int) bool {
iID := db.getRegistryID(filteredEntries[i])
jID := db.getRegistryID(filteredEntries[j])
return iID < jID
})
return filteredEntries
}
// matchesFilter checks if an entry matches the provided filter
func (db *MemoryDB) matchesFilter(entry *apiv0.ServerJSON, filter *ServerFilter) bool {
if filter == nil {
return true
}
// Check name filter
if filter.Name != nil && entry.Name != *filter.Name {
return false
}
// Check remote URL filter
if filter.RemoteURL != nil {
found := false
for _, remote := range entry.Remotes {
if remote.URL == *filter.RemoteURL {
found = true
break
}
}
if !found {
return false
}
}
return true
}
// getRegistryID safely extracts the registry ID from an entry
func (db *MemoryDB) getRegistryID(entry *apiv0.ServerJSON) string {
if entry.Meta != nil && entry.Meta.IOModelContextProtocolRegistry != nil {
return entry.Meta.IOModelContextProtocolRegistry.ID
}
return ""
}

View file

@ -5,9 +5,7 @@ import (
"encoding/json"
"errors"
"fmt"
"log"
"strings"
"time"
"github.com/google/uuid"
"github.com/jackc/pgx/v5"
@ -42,10 +40,9 @@ func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, erro
}, nil
}
// List retrieves ServerJSON entries with optional filtering and pagination
func (db *PostgreSQL) List(
ctx context.Context,
filter map[string]any,
filter *ServerFilter,
cursor string,
limit int,
) ([]*apiv0.ServerJSON, string, error) {
@ -57,30 +54,21 @@ func (db *PostgreSQL) List(
return nil, "", ctx.Err()
}
// Build WHERE clause for filtering - always filter by is_latest = true by default
// Build WHERE clause for filtering
var whereConditions []string
whereConditions = append(whereConditions, "(value->'_meta'->'io.modelcontextprotocol.registry'->>'is_latest')::boolean = true")
args := []any{}
argIndex := 1
// Add filters using JSON operators
for k, v := range filter {
switch k {
case "name":
if filter != nil {
if filter.Name != nil {
whereConditions = append(whereConditions, fmt.Sprintf("value->>'name' = $%d", argIndex))
args = append(args, v)
args = append(args, *filter.Name)
argIndex++
case "version":
whereConditions = append(whereConditions, fmt.Sprintf("value->'version_detail'->>'version' = $%d", argIndex))
args = append(args, v)
argIndex++
case "status":
whereConditions = append(whereConditions, fmt.Sprintf("value->>'status' = $%d", argIndex))
args = append(args, v)
argIndex++
case "remote_url":
}
if filter.RemoteURL != nil {
whereConditions = append(whereConditions, fmt.Sprintf("EXISTS (SELECT 1 FROM jsonb_array_elements(value->'remotes') AS remote WHERE remote->>'url' = $%d)", argIndex))
args = append(args, v)
args = append(args, *filter.RemoteURL)
argIndex++
}
}
@ -151,7 +139,6 @@ func (db *PostgreSQL) List(
return results, nextCursor, nil
}
// GetByID retrieves a single ServerJSON by its registry metadata ID
func (db *PostgreSQL) GetByID(ctx context.Context, id string) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
@ -183,27 +170,18 @@ func (db *PostgreSQL) GetByID(ctx context.Context, id string) (*apiv0.ServerJSON
return &serverJSON, nil
}
// Publish adds a new server to the simple servers table
func (db *PostgreSQL) Publish(ctx context.Context, serverDetail apiv0.ServerJSON, publisherExtensions map[string]interface{}, registryMetadata apiv0.RegistryExtensions) (*apiv0.ServerJSON, error) {
// CreateServer adds a new server to the database
func (db *PostgreSQL) CreateServer(ctx context.Context, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Create complete server with metadata
server := serverDetail // Copy the input
// Initialize meta if not present
if server.Meta == nil {
server.Meta = &apiv0.ServerMeta{}
// Get the ID from the registry metadata
if server.Meta == nil || server.Meta.IOModelContextProtocolRegistry == nil {
return nil, fmt.Errorf("server must have registry metadata with ID")
}
// Set registry metadata
server.Meta.IOModelContextProtocolRegistry = &registryMetadata
// Set publisher extensions if provided
if len(publisherExtensions) > 0 {
server.Meta.Publisher = publisherExtensions
}
id := server.Meta.IOModelContextProtocolRegistry.ID
// Marshal the complete server to JSONB
valueJSON, err := json.Marshal(server)
@ -217,130 +195,22 @@ func (db *PostgreSQL) Publish(ctx context.Context, serverDetail apiv0.ServerJSON
VALUES ($1, $2)
`
_, err = db.conn.Exec(ctx, query, registryMetadata.ID, valueJSON)
_, err = db.conn.Exec(ctx, query, id, valueJSON)
if err != nil {
return nil, fmt.Errorf("failed to insert server: %w", err)
}
return &server, nil
}
// ImportSeed imports initial data from a seed file into PostgreSQL
func (db *PostgreSQL) ImportSeed(ctx context.Context, seedFilePath string) error {
// Read seed data using the shared ReadSeedFile function
seedData, err := ReadSeedFile(ctx, seedFilePath)
if err != nil {
return fmt.Errorf("failed to read seed file: %w", err)
}
// Start a transaction for batch import
tx, err := db.conn.Begin(ctx)
if err != nil {
return fmt.Errorf("failed to start transaction: %w", err)
}
defer func() {
if rollbackErr := tx.Rollback(ctx); rollbackErr != nil && !errors.Is(rollbackErr, pgx.ErrTxClosed) {
log.Printf("Failed to rollback transaction: %v", rollbackErr)
}
}()
// Import each server
for _, record := range seedData {
// Marshal the server to JSONB
valueJSON, err := json.Marshal(record)
if err != nil {
return fmt.Errorf("failed to marshal server JSON for import: %w", err)
}
// Get registry ID from metadata
var registryID string
if record.Meta != nil && record.Meta.IOModelContextProtocolRegistry != nil {
registryID = record.Meta.IOModelContextProtocolRegistry.ID
} else {
registryID = uuid.New().String() // Generate ID if not present
}
// Insert into simple servers table
_, err = tx.Exec(ctx, "INSERT INTO servers (id, value) VALUES ($1, $2)", registryID, valueJSON)
if err != nil {
return fmt.Errorf("failed to import server %s: %w", record.Name, err)
}
}
// Commit the transaction
if err := tx.Commit(ctx); err != nil {
return fmt.Errorf("failed to commit seed import transaction: %w", err)
}
return nil
}
// UpdateLatestFlag updates the is_latest flag for a specific server record
func (db *PostgreSQL) UpdateLatestFlag(ctx context.Context, id string, isLatest bool) error {
if ctx.Err() != nil {
return ctx.Err()
}
query := `
UPDATE servers
SET value = jsonb_set(
jsonb_set(
value,
'{_meta,io.modelcontextprotocol.registry,is_latest}',
$1::jsonb
),
'{_meta,io.modelcontextprotocol.registry,updated_at}',
$2::jsonb
)
WHERE (value->'_meta'->'io.modelcontextprotocol.registry'->>'id') = $3
`
result, err := db.conn.Exec(ctx, query,
fmt.Sprintf("%t", isLatest),
fmt.Sprintf("\"%s\"", time.Now().Format(time.RFC3339)),
id)
if err != nil {
return fmt.Errorf("failed to update latest flag: %w", err)
}
if result.RowsAffected() == 0 {
return ErrNotFound
}
return nil
return server, nil
}
// UpdateServer updates an existing server record with new server details
func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, serverDetail apiv0.ServerJSON, publisherExtensions map[string]interface{}) (*apiv0.ServerJSON, error) {
func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Get the existing server to preserve registry metadata
existingServer, err := db.GetByID(ctx, id)
if err != nil {
return nil, err
}
// Update the server details while preserving registry metadata
updatedServer := serverDetail
if updatedServer.Meta == nil {
updatedServer.Meta = &apiv0.ServerMeta{}
}
// Preserve existing registry metadata but update timestamp
if existingServer.Meta != nil && existingServer.Meta.IOModelContextProtocolRegistry != nil {
updatedServer.Meta.IOModelContextProtocolRegistry = existingServer.Meta.IOModelContextProtocolRegistry
updatedServer.Meta.IOModelContextProtocolRegistry.UpdatedAt = time.Now()
}
// Update publisher extensions if provided
if len(publisherExtensions) > 0 {
updatedServer.Meta.Publisher = publisherExtensions
}
// Marshal updated server
valueJSON, err := json.Marshal(updatedServer)
valueJSON, err := json.Marshal(server)
if err != nil {
return nil, fmt.Errorf("failed to marshal updated server: %w", err)
}
@ -361,27 +231,10 @@ func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, serverDetail
return nil, ErrNotFound
}
return &updatedServer, nil
return server, 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,
}
}