Add db-backed implementation of RegistryService (#161)

This change adds a db-backed implementation of the `RegistryService`
interface.

There are multiple public interfaces added in this change,
specifically

* `internal/service` which exposes the `RegistryService` interface
  with an `Option` type used to configure filtering and pagination of
  SQL statements.
* `internal/service/db` which exposes means to create a
  `RegistryService` backed by a database.

Testing is implemented as integration tests running containers in the
same fashion as the `internal/db/sqlc` package. This might be revised
in the future by using mocks, but (a) a mock based implementation is
not going to be more succinct, and (b) does its job of covering most
of the new code.

Co-authored-by: Don Browne <[email protected]>

Author: blkt

go.mod

@@ -87,6 +87,7 @@ require (
 	github.com/ianlancetaylor/demangle v0.0.0-20250417193237-f615e6bd150b // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
+	github.com/jackc/puddle/v2 v2.2.2 // indirect
 	github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect

internal/api/extension/v0/routes.go

@@ -35,8 +35,10 @@ func Router(svc service.RegistryService) http.Handler {
 	r.Put("/registries/{registryName}", routes.upsertRegistry)
 	r.Delete("/registries/{registryName}", routes.deleteRegistry)
 
-	r.Put("/registries/{registryName}/servers/{serverName}/versions/{version}", routes.upsertVersion)
-	r.Delete("/registries/{registryName}/servers/{serverName}/versions/{version}", routes.deleteVersion)
+	r.Route("/registries/{registryName}/servers/{serverName}", func(r chi.Router) {
+		r.Put("/versions/{version}", routes.upsertVersion)
+		r.Delete("/versions/{version}", routes.deleteVersion)
+	})
 
 	return r
 }

internal/service/db/impl.go

@@ -0,0 +1,301 @@
+// Package database provides a database-backed implementation of the RegistryService interface
+package database
+
+import (
+	"context"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/google/uuid"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
+	upstreamv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
+	toolhivetypes "github.com/stacklok/toolhive/pkg/registry/registry"
+
+	"github.com/stacklok/toolhive-registry-server/internal/db/sqlc"
+	"github.com/stacklok/toolhive-registry-server/internal/service"
+)
+
+var (
+	// ErrBug is returned when a server is not found
+	ErrBug = errors.New("bug")
+)
+
+// options holds configuration options for the database service
+type options struct {
+	pool *pgxpool.Pool
+}
+
+// Option is a functional option for configuring the database service
+type Option func(*options) error
+
+// WithConnectionPool creates a new database-backed registry service with the
+// given pgx pool. The caller is responsible for closing the pool when it is
+// done.
+func WithConnectionPool(pool *pgxpool.Pool) Option {
+	return func(o *options) error {
+		if pool == nil {
+			return fmt.Errorf("pgx pool is required")
+		}
+		o.pool = pool
+		return nil
+	}
+}
+
+// dbService implements the RegistryService interface using a database backend
+type dbService struct {
+	pool *pgxpool.Pool
+}
+
+var _ service.RegistryService = (*dbService)(nil)
+
+// New creates a new database-backed registry service with the given options
+func New(opts ...Option) (service.RegistryService, error) {
+	o := &options{}
+
+	for _, opt := range opts {
+		if err := opt(o); err != nil {
+			return nil, err
+		}
+	}
+
+	return &dbService{
+		pool: o.pool,
+	}, nil
+}
+
+// CheckReadiness checks if the service is ready to serve requests
+func (s *dbService) CheckReadiness(ctx context.Context) error {
+	err := s.pool.Ping(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to ping database: %w", err)
+	}
+	return nil
+}
+
+// GetRegistry returns the registry data with metadata
+func (*dbService) GetRegistry(
+	_ context.Context,
+) (*toolhivetypes.UpstreamRegistry, string, error) {
+	return nil, "", service.ErrNotImplemented
+}
+
+// ListServers returns all servers in the registry
+func (s *dbService) ListServers(
+	ctx context.Context,
+	opts ...service.Option[service.ListServersOptions],
+) ([]*upstreamv0.ServerJSON, error) {
+	options := &service.ListServersOptions{}
+	for _, opt := range opts {
+		if err := opt(options); err != nil {
+			return nil, err
+		}
+	}
+
+	decoded, err := base64.StdEncoding.DecodeString(options.Cursor)
+	if err != nil {
+		return nil, fmt.Errorf("invalid cursor format: %w", err)
+	}
+	nextTime, err := time.Parse(time.RFC3339, string(decoded))
+	if err != nil {
+		return nil, fmt.Errorf("invalid cursor format: %w", err)
+	}
+
+	// Note: this function fetches a list of servers. In case no records are
+	// found, the called function should return an empty slice as it's
+	// customary in Go.
+	querierFunc := func(ctx context.Context, querier sqlc.Querier) ([]helper, error) {
+		servers, err := querier.ListServers(
+			ctx,
+			sqlc.ListServersParams{
+				Next: &nextTime,
+				Size: int64(options.Limit),
+			},
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		helpers := make([]helper, len(servers))
+		for i, server := range servers {
+			helpers[i] = listServersRowToHelper(server)
+		}
+
+		return helpers, nil
+	}
+
+	return s.sharedListServers(ctx, querierFunc)
+}
+
+// ListServerVersions implements RegistryService.ListServerVersions
+func (s *dbService) ListServerVersions(
+	ctx context.Context,
+	opts ...service.Option[service.ListServerVersionsOptions],
+) ([]*upstreamv0.ServerJSON, error) {
+	options := &service.ListServerVersionsOptions{}
+	for _, opt := range opts {
+		if err := opt(options); err != nil {
+			return nil, err
+		}
+	}
+
+	// Note: this function fetches a list of server versions. In case no records are
+	// found, the called function should return an empty slice as it's
+	// customary in Go.
+	querierFunc := func(ctx context.Context, querier sqlc.Querier) ([]helper, error) {
+		servers, err := querier.ListServerVersions(
+			ctx,
+			sqlc.ListServerVersionsParams{
+				Name: options.Name,
+				Next: options.Next,
+				Prev: options.Prev,
+				Size: int64(options.Limit),
+			},
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		helpers := make([]helper, len(servers))
+		for i, server := range servers {
+			helpers[i] = listServerVersionsRowToHelper(server)
+		}
+
+		return helpers, nil
+	}
+
+	return s.sharedListServers(ctx, querierFunc)
+}
+
+// GetServer returns a specific server by name
+func (s *dbService) GetServerVersion(
+	ctx context.Context,
+	opts ...service.Option[service.GetServerVersionOptions],
+) (*upstreamv0.ServerJSON, error) {
+	options := &service.GetServerVersionOptions{}
+	for _, opt := range opts {
+		if err := opt(options); err != nil {
+			return nil, err
+		}
+	}
+
+	// Note: this function fetches a single record given name and version.
+	// In case no record is found, the called function should return an
+	// `sql.ErrNoRows` error as it's customary in Go.
+	querierFunc := func(ctx context.Context, querier sqlc.Querier) ([]helper, error) {
+		server, err := querier.GetServerVersion(
+			ctx,
+			sqlc.GetServerVersionParams{
+				Name:    options.Name,
+				Version: options.Version,
+			},
+		)
+		if err != nil {
+			return nil, err
+		}
+
+		return []helper{getServerVersionRowToHelper(server)}, nil
+	}
+
+	res, err := s.sharedListServers(ctx, querierFunc)
+	if err != nil {
+		return nil, err
+	}
+
+	// Note: the `queryFunc` function is expected to return an error
+	// sooner if no records are found, so getting this far with
+	// a length result slice other than 1 means there's a bug.
+	if len(res) != 1 {
+		return nil, fmt.Errorf("%w: number of servers returned is not 1", ErrBug)
+	}
+
+	return res[0], nil
+}
+
+// querierFunction is a function that uses the given querier object to run the
+// main extraction. As of the time of this writing, its main use is accessing
+// the `mcp_server` table in a type-agnostic way. This is to overcome a
+// limitation of sqlc that prevents us from having the exact same go type
+// despite the fact that the underlying columns returned are the same.
+//
+// Note that the underlying table does not have to be the `mcp_server` table
+// as it is used now, as long as the result is a slice of helpers.
+type querierFunction func(ctx context.Context, querier sqlc.Querier) ([]helper, error)
+
+// sharedListServers is a helper function to list servers and mapping them to
+// the API schema.
+//
+// Its responsibilities are:
+// * Begin a transaction
+// * Execute the querier function
+// * List packages and remotes using the server IDs
+// * Map the results to the API schema
+// * Return the results
+//
+// The argument `querierFunc` is a function that uses the given querier object
+// to run the main extraction. Note that the underlying table does not have
+// to be the `mcp_server` table as it is used now, as long as the result is a
+// slice of helpers.
+func (s *dbService) sharedListServers(
+	ctx context.Context,
+	querierFunc querierFunction,
+) ([]*upstreamv0.ServerJSON, error) {
+	tx, err := s.pool.BeginTx(ctx, pgx.TxOptions{
+		IsoLevel:   pgx.ReadCommitted,
+		AccessMode: pgx.ReadOnly,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("failed to begin transaction: %w", err)
+	}
+	defer func() {
+		err := tx.Rollback(ctx)
+		if err != nil && !errors.Is(err, pgx.ErrTxClosed) {
+			// TODO: log the rollback error (add proper logging)
+			_ = err
+		}
+	}()
+
+	querier := sqlc.New(tx)
+
+	servers, err := querierFunc(ctx, querier)
+	if err != nil {
+		return nil, err
+	}
+
+	ids := make([]uuid.UUID, len(servers))
+	for i, server := range servers {
+		ids[i] = server.ID
+	}
+
+	packages, err := querier.ListServerPackages(ctx, ids)
+	if err != nil {
+		return nil, err
+	}
+	packagesMap := make(map[uuid.UUID][]sqlc.McpServerPackage)
+	for _, pkg := range packages {
+		packagesMap[pkg.ServerID] = append(packagesMap[pkg.ServerID], pkg)
+	}
+
+	remotes, err := querier.ListServerRemotes(ctx, ids)
+	if err != nil {
+		return nil, err
+	}
+	remotesMap := make(map[uuid.UUID][]sqlc.McpServerRemote)
+	for _, remote := range remotes {
+		remotesMap[remote.ServerID] = append(remotesMap[remote.ServerID], remote)
+	}
+
+	result := make([]*upstreamv0.ServerJSON, 0, len(servers))
+	for _, dbServer := range servers {
+		server := helperToServer(
+			dbServer,
+			packagesMap[dbServer.ID],
+			remotesMap[dbServer.ID],
+		)
+		result = append(result, &server)
+	}
+
+	return result, nil
+}