package database import ( "context" "encoding/json" "errors" "fmt" "strings" "github.com/google/uuid" "github.com/jackc/pgx/v5" apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0" ) // 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 } func (db *PostgreSQL) List( ctx context.Context, filter *ServerFilter, cursor string, limit int, ) ([]*apiv0.ServerJSON, string, error) { if limit <= 0 { limit = 10 } if ctx.Err() != nil { return nil, "", ctx.Err() } // Build WHERE clause for filtering var whereConditions []string args := []any{} argIndex := 1 // Add filters using JSON operators if filter != nil { if filter.Name != nil { whereConditions = append(whereConditions, fmt.Sprintf("value->>'name' = $%d", argIndex)) args = append(args, *filter.Name) argIndex++ } 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, *filter.RemoteURL) argIndex++ } } // Add cursor pagination using registry metadata ID if cursor != "" { if _, err := uuid.Parse(cursor); err != nil { return nil, "", fmt.Errorf("invalid cursor format: %w", err) } whereConditions = append(whereConditions, fmt.Sprintf("(value->'_meta'->'io.modelcontextprotocol.registry/official'->>'id') > $%d", argIndex)) args = append(args, cursor) argIndex++ } // Build the WHERE clause whereClause := "" if len(whereConditions) > 0 { whereClause = "WHERE " + strings.Join(whereConditions, " AND ") } // Simple query on servers table query := fmt.Sprintf(` SELECT value FROM servers %s ORDER BY (value->'_meta'->'io.modelcontextprotocol.registry/official'->>'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 []*apiv0.ServerJSON for rows.Next() { var valueJSON []byte err := rows.Scan(&valueJSON) if err != nil { return nil, "", fmt.Errorf("failed to scan server row: %w", err) } // Parse the complete ServerJSON from JSONB var serverJSON apiv0.ServerJSON if err := json.Unmarshal(valueJSON, &serverJSON); err != nil { return nil, "", fmt.Errorf("failed to unmarshal server JSON: %w", err) } results = append(results, &serverJSON) } if err := rows.Err(); err != nil { return nil, "", fmt.Errorf("error iterating rows: %w", err) } // Determine next cursor using registry metadata ID nextCursor := "" if len(results) > 0 && len(results) >= limit { lastResult := results[len(results)-1] if lastResult.Meta != nil && lastResult.Meta.Official != nil { nextCursor = lastResult.Meta.Official.ID } } return results, nextCursor, nil } func (db *PostgreSQL) GetByID(ctx context.Context, id string) (*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } query := ` SELECT value FROM servers WHERE (value->'_meta'->'io.modelcontextprotocol.registry/official'->>'id') = $1 ` var valueJSON []byte err := db.conn.QueryRow(ctx, query, id).Scan(&valueJSON) 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 the complete ServerJSON from JSONB var serverJSON apiv0.ServerJSON if err := json.Unmarshal(valueJSON, &serverJSON); err != nil { return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err) } return &serverJSON, nil } // 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() } // Get the ID from the registry metadata if server.Meta == nil || server.Meta.Official == nil { return nil, fmt.Errorf("server must have registry metadata with ID") } id := server.Meta.Official.ID // Marshal the complete server to JSONB valueJSON, err := json.Marshal(server) if err != nil { return nil, fmt.Errorf("failed to marshal server JSON: %w", err) } // Insert into simple servers table query := ` INSERT INTO servers (id, value) VALUES ($1, $2) ` _, err = db.conn.Exec(ctx, query, id, valueJSON) if err != nil { return nil, fmt.Errorf("failed to insert server: %w", err) } return server, nil } // UpdateServer updates an existing server record with new server details func (db *PostgreSQL) UpdateServer(ctx context.Context, id string, server *apiv0.ServerJSON) (*apiv0.ServerJSON, error) { if ctx.Err() != nil { return nil, ctx.Err() } // Marshal updated server valueJSON, err := json.Marshal(server) if err != nil { return nil, fmt.Errorf("failed to marshal updated server: %w", err) } // Update the complete server record in simple table query := ` UPDATE servers SET value = $1 WHERE (value->'_meta'->'io.modelcontextprotocol.registry/official'->>'id') = $2 ` result, err := db.conn.Exec(ctx, query, valueJSON, id) if err != nil { return nil, fmt.Errorf("failed to update server: %w", err) } if result.RowsAffected() == 0 { return nil, ErrNotFound } return server, nil } // Close closes the database connection func (db *PostgreSQL) Close() error { return db.conn.Close(context.Background()) }