registry/internal/database/postgres.go
Radoslav Dimitrov 736d376858
fix(db): raise pgxpool MaxConns + PG max_connections + set PG resources (#1221)
## Summary

Follow-up to #1215 and #1220. Both of those addressed *individual* slow
queries; this PR addresses *concurrency* — under today's scraper load
the cursor query is fast on average (mean 42ms) but the connection pool
saturates and queue depth blows up at the Go HTTP layer.

## Diagnosis

`pg_stat_statements` (added in #1215) made this possible to see:

```
max_ms   mean_ms  calls    template
10,755     41.8   94,009   plain cursor pagination, no filter
 7,277    214.7   14,770   ILIKE substring filter
 7,638    330.8    7,656   ILIKE substring + is_latest filter
```

Mean times are healthy. But `max_exec_time` of 7–10s on the cursor
query, combined with sustained ~15 req/s from scrapers (ServiceNow's
148.139.x.x range, anonymous `node` user-agent, etc) saturates
`MaxConns=30 × 2 pods = 60`. New requests queue at the Go HTTP layer;
nginx-side p99 hits 35s; eventually scrapers time out at 60s, retry, and
amplify the queue.

Today's two ongoing alerts are both this pattern:
- 17:35–17:40 UTC `Availability dropped below 95%` — 4,526 GET requests
in 5 min, hundreds of 504s
- 18:17 UTC `Publish Endpoint Latency` re-fire — same scraper
concurrency dragging the publish path

Critically, the symptoms today were also visible during yesterday's
incident, but yesterday's broken cursor (#1215) was the dominant cause.
After #1215 the cursor is fast individually; concurrency now becomes the
next bottleneck.

## Changes

### pgxpool (`internal/database/postgres.go`)

| Setting | Before | After |
|---------|-------:|------:|
| `MaxConns` | 30 | **60** per pod |
| `MinConns` | 5 | **10** per pod |
| `MaxConnIdleTime` | 30 min | (unchanged) |
| `MaxConnLifetime` | 2 h | (unchanged) |

2 pods × 60 = 120 total app connections. Cuts queue depth roughly in
half at current scraper load.

### PG cluster (`deploy/pkg/k8s/postgres.go`)

- `max_connections: 100 → 200` — required to support the larger pool.
120 (app) + ~10 (PG internals: autovacuum, replication, admin) + 70
headroom.
- Added explicit `resources:` block — previously the pod had no resource
limits, making node-level OOM behaviour unpredictable.

| | Request | Limit |
|---|--------:|------:|
| memory | 512Mi | **4Gi** |
| cpu | 200m | 1500m |

## Resource budget

| Node | Now | After |
|------|----:|------:|
| dy89 (PG node) | 39% mem (~2.4 GiB / 6 GiB) | ~65% mem worst-case
(~3.9 GiB) |
| 2yxm | 36% mem | unchanged |

Both nodes well within capacity. CPU usage <50% on both, plenty of
headroom.

PG worst-case memory math:
- 200 conns × ~15 MB per backend = 3 GiB
- + `shared_buffers` 128 MiB
- + `maintenance_work_mem`, OS overhead, etc.
- ≈ 3–4 GiB total

## Deployment caveat

`max_connections` is a postmaster-level setting → CNPG triggers a PG
restart on the change. Same in-pod restart shape as the
pg_stat_statements deploy yesterday — registry pods see ~30s of DB
unavailability, covered by v1.7.1's retry-with-backoff (8 attempts, 1→8s
capped). One registry pod may bounce once before recovering, like
yesterday.

**Time the merge for a low-traffic UTC window.**

Order of operations matters within the deploy itself:
1. Pulumi applies the new CNPG spec → PG restarts with
`max_connections=200`
2. Rolling deploy of registry pods picks up `MaxConns=60` config
3. New conns are accepted up to 200 limit

Pulumi's standard ordering does step 1 before step 2 in this scenario
(Pulumi resource graph: CNPG cluster precedes Deployment). If for any
reason it doesn't, the worst case is `too many connections` errors
during a small window — Self-correcting once the rollout completes.

## Test plan

- [x] `go build ./...` clean for app + deploy
- [x] `make lint` clean
- [x] `go test -race ./internal/database/...` green
- [ ] On merge: staging deploy applies the spec change cleanly; PG
restarts; pool reaches 60 max
- [ ] On prod deploy: same; verify no `too many connections` errors
during the rollout window
- [ ] After deploy: query `SHOW max_connections` returns 200; query
pg_stat_statements after a scraper burst, confirm `max_exec_time` for
cursor query no longer hits 10s

## Out of scope

- **ILIKE substring search** (`server_name ILIKE '%foo%'`) is
unindexable. Three pg_stat_statements variants run with means 141–333ms
and max 4–7s. Worth replacing with a `pg_trgm` GIN index or full-text
search in a separate PR.
- **Per-IP nginx rate limiting** — defends against scraper retry storms
regardless of pool size.
- **`Cache-Control` on `/v0/servers`** — let nginx absorb
scraper-repeated cursors.
- **Pre-existing `superfluous WriteHeader` warnings** — separate Huma
framework issue.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2026-04-28 15:21:24 -04:00

938 lines
29 KiB
Go

package database
import (
"context"
"encoding/json"
"errors"
"fmt"
"log"
"strings"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgconn"
"github.com/jackc/pgx/v5/pgxpool"
apiv0 "github.com/modelcontextprotocol/registry/pkg/api/v0"
"github.com/modelcontextprotocol/registry/pkg/model"
)
// PostgreSQL is an implementation of the Database interface using PostgreSQL
type PostgreSQL struct {
pool *pgxpool.Pool
}
// Executor is an interface for executing queries (satisfied by both pgx.Tx and pgxpool.Pool)
type Executor interface {
Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}
// getExecutor returns the appropriate executor (transaction or pool)
func (db *PostgreSQL) getExecutor(tx pgx.Tx) Executor {
if tx != nil {
return tx
}
return db.pool
}
// NewPostgreSQL creates a new instance of the PostgreSQL database
func NewPostgreSQL(ctx context.Context, connectionURI string) (*PostgreSQL, error) {
// Parse connection config for pool settings
config, err := pgxpool.ParseConfig(connectionURI)
if err != nil {
return nil, fmt.Errorf("failed to parse PostgreSQL config: %w", err)
}
// Configure pool for stability-focused defaults.
// MaxConns was 30 per pod (60 total across 2 replicas) which saturated under
// the 2026-04-28 scraper bursts (15 req/s on /v0/servers caused queue blowup
// even though individual queries were fast). 60 per pod gives 120 total,
// leaving 80 of PG max_connections=200 for autovacuum/admin/headroom.
config.MaxConns = 60 // Handle scraper-burst concurrent load
config.MinConns = 10 // 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 = pool.Ping(ctx); err != nil {
return nil, fmt.Errorf("failed to ping PostgreSQL: %w", err)
}
// 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{
pool: pool,
}, nil
}
// buildFilterConditions constructs WHERE clause conditions from a ServerFilter
//
//nolint:unparam // argIndex is always 1 currently but kept for API flexibility
func buildFilterConditions(filter *ServerFilter, argIndex int) ([]string, []any, int) {
var conditions []string
var args []any
if filter == nil {
return conditions, args, argIndex
}
if filter.Name != nil {
conditions = append(conditions, fmt.Sprintf("server_name = $%d", argIndex))
args = append(args, *filter.Name)
argIndex++
}
if filter.RemoteURL != nil {
// Use a JSONB containment predicate so the planner can use the GIN index
// idx_servers_json_remotes on (value -> 'remotes'). The previously-used
// EXISTS / jsonb_array_elements / ->> form is logically equivalent but
// the planner can't translate the per-row array unfolding into a GIN
// search — it falls back to scanning every row. validateNoDuplicateRemoteURLs
// in the publish path ran this filter and was measured at ~10s on prod's
// 21K-row table on 2026-04-28; the containment form is GIN-indexable and
// completes in low single-digit ms.
conditions = append(conditions, fmt.Sprintf("value -> 'remotes' @> jsonb_build_array(jsonb_build_object('url', $%d::text))", argIndex))
args = append(args, *filter.RemoteURL)
argIndex++
}
if filter.UpdatedSince != nil {
conditions = append(conditions, fmt.Sprintf("updated_at > $%d", argIndex))
args = append(args, *filter.UpdatedSince)
argIndex++
}
if filter.SubstringName != nil {
conditions = append(conditions, fmt.Sprintf("server_name ILIKE $%d", argIndex))
args = append(args, "%"+*filter.SubstringName+"%")
argIndex++
}
if filter.Version != nil {
conditions = append(conditions, fmt.Sprintf("version = $%d", argIndex))
args = append(args, *filter.Version)
argIndex++
}
if filter.IsLatest != nil {
conditions = append(conditions, fmt.Sprintf("is_latest = $%d", argIndex))
args = append(args, *filter.IsLatest)
argIndex++
}
if filter.IncludeDeleted == nil || !*filter.IncludeDeleted {
conditions = append(conditions, "status != 'deleted'")
}
return conditions, args, argIndex
}
// addCursorCondition adds pagination cursor condition to WHERE clause.
//
// The compound cursor uses a row-constructor comparison so PostgreSQL can seek
// directly into the (server_name, version) B-tree index. The OR-decomposed form
// `server_name > X OR (server_name = X AND version > Y)` is logically equivalent
// but PostgreSQL's planner cannot use it for an index seek — it scans the index
// from the start and filters everything before the cursor, making cost grow
// linearly with cursor depth (a 20K-row table at a deep cursor took ~760ms in
// prod). The row-constructor form `(server_name, version) > (X, Y)` is special-
// cased and stays constant-time regardless of cursor depth.
func addCursorCondition(cursor string, argIndex int) (string, []any, int) {
if cursor == "" {
return "", nil, argIndex
}
// Parse cursor format: "serverName:version"
parts := strings.SplitN(cursor, ":", 2)
if len(parts) == 2 {
cursorServerName := parts[0]
cursorVersion := parts[1]
condition := fmt.Sprintf("(server_name, version) > ($%d, $%d)", argIndex, argIndex+1)
return condition, []any{cursorServerName, cursorVersion}, argIndex + 2
}
// Fallback for malformed cursor - treat as server name only for backwards compatibility
condition := fmt.Sprintf("server_name > $%d", argIndex)
return condition, []any{cursor}, argIndex + 1
}
func (db *PostgreSQL) ListServers(
ctx context.Context,
tx pgx.Tx,
filter *ServerFilter,
cursor string,
limit int,
) ([]*apiv0.ServerResponse, string, error) {
if limit <= 0 {
limit = 10
}
if ctx.Err() != nil {
return nil, "", ctx.Err()
}
// Build WHERE clause conditions
argIndex := 1
whereConditions, args, argIndex := buildFilterConditions(filter, argIndex)
// Add cursor pagination
cursorCondition, cursorArgs, argIndex := addCursorCondition(cursor, argIndex)
if cursorCondition != "" {
whereConditions = append(whereConditions, cursorCondition)
args = append(args, cursorArgs...)
}
_ = argIndex // Silence unused variable warning
// Build the WHERE clause
whereClause := ""
if len(whereConditions) > 0 {
whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
}
// Query servers table with hybrid column/JSON data
query := fmt.Sprintf(`
SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
FROM servers
%s
ORDER BY server_name, version
LIMIT $%d
`, whereClause, argIndex)
args = append(args, limit)
rows, err := db.getExecutor(tx).Query(ctx, query, args...)
if err != nil {
return nil, "", fmt.Errorf("failed to query servers: %w", err)
}
defer rows.Close()
var results []*apiv0.ServerResponse
for rows.Next() {
var serverName, version, status string
var statusChangedAt, publishedAt, updatedAt time.Time
var statusMessage *string
var isLatest bool
var valueJSON []byte
err := rows.Scan(&serverName, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
if err != nil {
return nil, "", fmt.Errorf("failed to scan server row: %w", err)
}
// Parse the 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)
}
// Build ServerResponse with separated metadata
serverResponse := &apiv0.ServerResponse{
Server: serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(status),
StatusChangedAt: statusChangedAt,
StatusMessage: statusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
results = append(results, serverResponse)
}
if err := rows.Err(); err != nil {
return nil, "", fmt.Errorf("error iterating rows: %w", err)
}
// Determine next cursor using compound serverName:version format
nextCursor := ""
if len(results) > 0 && len(results) >= limit {
lastResult := results[len(results)-1]
nextCursor = lastResult.Server.Name + ":" + lastResult.Server.Version
}
return results, nextCursor, nil
}
// GetServerByName retrieves the latest version of a server by server name
func (db *PostgreSQL) GetServerByName(ctx context.Context, tx pgx.Tx, serverName string, includeDeleted bool) (*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Build filter conditions
isLatest := true
filter := &ServerFilter{
Name: &serverName,
IsLatest: &isLatest,
IncludeDeleted: &includeDeleted,
}
argIndex := 1
whereConditions, args, _ := buildFilterConditions(filter, argIndex)
whereClause := ""
if len(whereConditions) > 0 {
whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
}
query := fmt.Sprintf(`
SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
FROM servers
%s
ORDER BY published_at DESC
LIMIT 1
`, whereClause)
var name, version, status string
var statusChangedAt, publishedAt, updatedAt time.Time
var statusMessage *string
var valueJSON []byte
err := db.getExecutor(tx).QueryRow(ctx, query, args...).Scan(&name, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, ErrNotFound
}
return nil, fmt.Errorf("failed to get server by name: %w", err)
}
// Parse the 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)
}
// Build ServerResponse with separated metadata
serverResponse := &apiv0.ServerResponse{
Server: serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(status),
StatusChangedAt: statusChangedAt,
StatusMessage: statusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
return serverResponse, nil
}
// GetServerByNameAndVersion retrieves a specific version of a server by server name and version
func (db *PostgreSQL) GetServerByNameAndVersion(ctx context.Context, tx pgx.Tx, serverName string, version string, includeDeleted bool) (*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Build filter conditions
filter := &ServerFilter{
Name: &serverName,
Version: &version,
IncludeDeleted: &includeDeleted,
}
argIndex := 1
whereConditions, args, _ := buildFilterConditions(filter, argIndex)
whereClause := ""
if len(whereConditions) > 0 {
whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
}
query := fmt.Sprintf(`
SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
FROM servers
%s
LIMIT 1
`, whereClause)
var name, vers, status string
var statusChangedAt, publishedAt, updatedAt time.Time
var statusMessage *string
var isLatest bool
var valueJSON []byte
err := db.getExecutor(tx).QueryRow(ctx, query, args...).Scan(&name, &vers, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, ErrNotFound
}
return nil, fmt.Errorf("failed to get server by name and version: %w", err)
}
// Parse the 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)
}
// Build ServerResponse with separated metadata
serverResponse := &apiv0.ServerResponse{
Server: serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(status),
StatusChangedAt: statusChangedAt,
StatusMessage: statusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
return serverResponse, nil
}
// GetAllVersionsByServerName retrieves all versions of a server by server name
func (db *PostgreSQL) GetAllVersionsByServerName(ctx context.Context, tx pgx.Tx, serverName string, includeDeleted bool) ([]*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Build filter conditions
filter := &ServerFilter{
Name: &serverName,
IncludeDeleted: &includeDeleted,
}
argIndex := 1
whereConditions, args, _ := buildFilterConditions(filter, argIndex)
whereClause := ""
if len(whereConditions) > 0 {
whereClause = "WHERE " + strings.Join(whereConditions, " AND ")
}
query := fmt.Sprintf(`
SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
FROM servers
%s
ORDER BY published_at DESC
`, whereClause)
rows, err := db.getExecutor(tx).Query(ctx, query, args...)
if err != nil {
return nil, fmt.Errorf("failed to query server versions: %w", err)
}
defer rows.Close()
var results []*apiv0.ServerResponse
for rows.Next() {
var name, version, status string
var statusChangedAt, publishedAt, updatedAt time.Time
var statusMessage *string
var isLatest bool
var valueJSON []byte
err := rows.Scan(&name, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &valueJSON)
if err != nil {
return nil, fmt.Errorf("failed to scan server row: %w", err)
}
// Parse the 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)
}
// Build ServerResponse with separated metadata
serverResponse := &apiv0.ServerResponse{
Server: serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(status),
StatusChangedAt: statusChangedAt,
StatusMessage: statusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
results = append(results, serverResponse)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("error iterating rows: %w", err)
}
if len(results) == 0 {
return nil, ErrNotFound
}
return results, nil
}
// CreateServer inserts a new server version with official metadata
func (db *PostgreSQL) CreateServer(ctx context.Context, tx pgx.Tx, serverJSON *apiv0.ServerJSON, officialMeta *apiv0.RegistryExtensions) (*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Validate inputs
if serverJSON == nil || officialMeta == nil {
return nil, fmt.Errorf("serverJSON and officialMeta are required")
}
if serverJSON.Name == "" || serverJSON.Version == "" {
return nil, fmt.Errorf("server name and version are required")
}
// Marshal the ServerJSON to JSONB
valueJSON, err := json.Marshal(serverJSON)
if err != nil {
return nil, fmt.Errorf("failed to marshal server JSON: %w", err)
}
// Insert the new server version using composite primary key
insertQuery := `
INSERT INTO servers (server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
`
_, err = db.getExecutor(tx).Exec(ctx, insertQuery,
serverJSON.Name,
serverJSON.Version,
string(officialMeta.Status),
officialMeta.StatusChangedAt,
officialMeta.StatusMessage,
officialMeta.PublishedAt,
officialMeta.UpdatedAt,
officialMeta.IsLatest,
valueJSON,
)
if err != nil {
return nil, fmt.Errorf("failed to insert server: %w", err)
}
// Return the complete ServerResponse
serverResponse := &apiv0.ServerResponse{
Server: *serverJSON,
Meta: apiv0.ResponseMeta{
Official: officialMeta,
},
}
return serverResponse, nil
}
// UpdateServer updates an existing server record with new server details
func (db *PostgreSQL) UpdateServer(ctx context.Context, tx pgx.Tx, serverName, version string, serverJSON *apiv0.ServerJSON) (*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Validate inputs
if serverJSON == nil {
return nil, fmt.Errorf("serverJSON is required")
}
// Ensure the serverJSON matches the provided serverName and version
if serverJSON.Name != serverName || serverJSON.Version != version {
return nil, fmt.Errorf("%w: server name and version in JSON must match parameters", ErrInvalidInput)
}
// Marshal updated ServerJSON
valueJSON, err := json.Marshal(serverJSON)
if err != nil {
return nil, fmt.Errorf("failed to marshal updated server: %w", err)
}
// Update only the JSON data (keep existing metadata columns)
query := `
UPDATE servers
SET value = $1, updated_at = NOW()
WHERE server_name = $2 AND version = $3
RETURNING server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest
`
var name, vers, status string
var statusChangedAt, publishedAt, updatedAt time.Time
var statusMessage *string
var isLatest bool
err = db.getExecutor(tx).QueryRow(ctx, query, valueJSON, serverName, version).Scan(&name, &vers, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, ErrNotFound
}
return nil, fmt.Errorf("failed to update server: %w", err)
}
// Return the updated ServerResponse
serverResponse := &apiv0.ServerResponse{
Server: *serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(status),
StatusChangedAt: statusChangedAt,
StatusMessage: statusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
return serverResponse, nil
}
// SetServerStatus updates the status of a specific server version
func (db *PostgreSQL) SetServerStatus(ctx context.Context, tx pgx.Tx, serverName, version string, status model.Status, statusMessage *string) (*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Update the status and related fields
// Only update status_changed_at when status actually changes
query := `
UPDATE servers
SET
status = $1,
status_changed_at = CASE WHEN status != $1::varchar THEN NOW() ELSE status_changed_at END,
updated_at = NOW(),
status_message = $4
WHERE server_name = $2 AND version = $3
RETURNING server_name, version, status, value, published_at, updated_at, is_latest, status_changed_at, status_message
`
var name, vers, currentStatus string
var publishedAt, updatedAt, statusChangedAt time.Time
var isLatest bool
var valueJSON []byte
var resultStatusMessage *string
err := db.getExecutor(tx).QueryRow(ctx, query, string(status), serverName, version, statusMessage).Scan(&name, &vers, &currentStatus, &valueJSON, &publishedAt, &updatedAt, &isLatest, &statusChangedAt, &resultStatusMessage)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, ErrNotFound
}
return nil, fmt.Errorf("failed to update server status: %w", err)
}
// Unmarshal the JSON data
var serverJSON apiv0.ServerJSON
if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
}
// Return the updated ServerResponse
serverResponse := &apiv0.ServerResponse{
Server: serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(currentStatus),
StatusChangedAt: statusChangedAt,
StatusMessage: resultStatusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
return serverResponse, nil
}
// SetAllVersionsStatus updates the status of all versions of a server in a single query
func (db *PostgreSQL) SetAllVersionsStatus(ctx context.Context, tx pgx.Tx, serverName string, status model.Status, statusMessage *string) ([]*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
// Update the status and related fields for all versions
// Only update rows where status or status_message actually changes
// Only update status_changed_at when status actually changes
query := `
UPDATE servers
SET
status = $1,
status_changed_at = CASE WHEN status != $1::varchar THEN NOW() ELSE status_changed_at END,
updated_at = NOW(),
status_message = $2
WHERE server_name = $3
AND (status != $1::varchar OR status_message IS DISTINCT FROM $2)
RETURNING server_name, version, status, value, published_at, updated_at, is_latest, status_changed_at, status_message
`
rows, err := db.getExecutor(tx).Query(ctx, query, string(status), statusMessage, serverName)
if err != nil {
return nil, fmt.Errorf("failed to update all server versions status: %w", err)
}
defer rows.Close()
var results []*apiv0.ServerResponse
for rows.Next() {
var name, vers, currentStatus string
var publishedAt, updatedAt, statusChangedAt time.Time
var isLatest bool
var valueJSON []byte
var resultStatusMessage *string
if err := rows.Scan(&name, &vers, &currentStatus, &valueJSON, &publishedAt, &updatedAt, &isLatest, &statusChangedAt, &resultStatusMessage); err != nil {
return nil, fmt.Errorf("failed to scan server row: %w", err)
}
// Unmarshal the JSON data
var serverJSON apiv0.ServerJSON
if err := json.Unmarshal(valueJSON, &serverJSON); err != nil {
return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
}
serverResponse := &apiv0.ServerResponse{
Server: serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(currentStatus),
StatusChangedAt: statusChangedAt,
StatusMessage: resultStatusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
results = append(results, serverResponse)
}
if err := rows.Err(); err != nil {
return nil, fmt.Errorf("error iterating server rows: %w", err)
}
if len(results) == 0 {
return nil, ErrNotFound
}
return results, nil
}
// InTransaction executes a function within a database transaction
func (db *PostgreSQL) InTransaction(ctx context.Context, fn func(ctx context.Context, tx pgx.Tx) error) error {
if ctx.Err() != nil {
return ctx.Err()
}
tx, err := db.pool.Begin(ctx)
if err != nil {
return fmt.Errorf("failed to begin transaction: %w", err)
}
//nolint:contextcheck // Intentionally using separate context for rollback to ensure cleanup even if request is cancelled
defer func() {
rollbackCtx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
defer cancel()
if rbErr := tx.Rollback(rollbackCtx); rbErr != nil && !errors.Is(rbErr, pgx.ErrTxClosed) {
log.Printf("failed to rollback transaction: %v", rbErr)
}
}()
if err := fn(ctx, tx); err != nil {
return err
}
if err := tx.Commit(ctx); err != nil {
return fmt.Errorf("failed to commit transaction: %w", err)
}
return nil
}
// AcquirePublishLock acquires an exclusive advisory lock for publishing a server
// This prevents race conditions when multiple versions are published concurrently
// Using pg_advisory_xact_lock which auto-releases on transaction end
func (db *PostgreSQL) AcquirePublishLock(ctx context.Context, tx pgx.Tx, serverName string) error {
if ctx.Err() != nil {
return ctx.Err()
}
lockID := hashServerName(serverName)
if _, err := db.getExecutor(tx).Exec(ctx, "SELECT pg_advisory_xact_lock($1)", lockID); err != nil {
return fmt.Errorf("failed to acquire publish lock: %w", err)
}
return nil
}
// hashServerName creates a consistent hash of the server name for advisory locking
// We use FNV-1a hash and mask to 63 bits to fit in PostgreSQL's bigint range
func hashServerName(name string) int64 {
const (
offset64 = 14695981039346656037
prime64 = 1099511628211
)
hash := uint64(offset64)
for i := 0; i < len(name); i++ {
hash ^= uint64(name[i])
hash *= prime64
}
return int64(hash & 0x7FFFFFFFFFFFFFFF)
}
// GetCurrentLatestVersion retrieves the current latest version of a server by server name
func (db *PostgreSQL) GetCurrentLatestVersion(ctx context.Context, tx pgx.Tx, serverName string) (*apiv0.ServerResponse, error) {
if ctx.Err() != nil {
return nil, ctx.Err()
}
executor := db.getExecutor(tx)
query := `
SELECT server_name, version, status, status_changed_at, status_message, published_at, updated_at, is_latest, value
FROM servers
WHERE server_name = $1 AND is_latest = true
`
row := executor.QueryRow(ctx, query, serverName)
var name, version, status string
var statusChangedAt, publishedAt, updatedAt time.Time
var statusMessage *string
var isLatest bool
var jsonValue []byte
err := row.Scan(&name, &version, &status, &statusChangedAt, &statusMessage, &publishedAt, &updatedAt, &isLatest, &jsonValue)
if err != nil {
if errors.Is(err, pgx.ErrNoRows) {
return nil, ErrNotFound
}
return nil, fmt.Errorf("failed to scan server row: %w", err)
}
// Parse the JSON value to get the server details
var serverJSON apiv0.ServerJSON
if err := json.Unmarshal(jsonValue, &serverJSON); err != nil {
return nil, fmt.Errorf("failed to unmarshal server JSON: %w", err)
}
// Build ServerResponse with separated metadata
serverResponse := &apiv0.ServerResponse{
Server: serverJSON,
Meta: apiv0.ResponseMeta{
Official: &apiv0.RegistryExtensions{
Status: model.Status(status),
StatusChangedAt: statusChangedAt,
StatusMessage: statusMessage,
PublishedAt: publishedAt,
UpdatedAt: updatedAt,
IsLatest: isLatest,
},
},
}
return serverResponse, nil
}
// CountServerVersions counts the number of versions for a server
func (db *PostgreSQL) CountServerVersions(ctx context.Context, tx pgx.Tx, serverName string) (int, error) {
if ctx.Err() != nil {
return 0, ctx.Err()
}
executor := db.getExecutor(tx)
query := `SELECT COUNT(*) FROM servers WHERE server_name = $1`
var count int
err := executor.QueryRow(ctx, query, serverName).Scan(&count)
if err != nil {
return 0, fmt.Errorf("failed to count server versions: %w", err)
}
return count, nil
}
// CheckVersionExists checks if a specific version exists for a server
func (db *PostgreSQL) CheckVersionExists(ctx context.Context, tx pgx.Tx, serverName, version string) (bool, error) {
if ctx.Err() != nil {
return false, ctx.Err()
}
executor := db.getExecutor(tx)
query := `SELECT EXISTS(SELECT 1 FROM servers WHERE server_name = $1 AND version = $2)`
var exists bool
err := executor.QueryRow(ctx, query, serverName, version).Scan(&exists)
if err != nil {
return false, fmt.Errorf("failed to check version existence: %w", err)
}
return exists, nil
}
// UnmarkAsLatest marks the current latest version of a server as no longer latest
func (db *PostgreSQL) UnmarkAsLatest(ctx context.Context, tx pgx.Tx, serverName string) error {
if ctx.Err() != nil {
return ctx.Err()
}
executor := db.getExecutor(tx)
query := `UPDATE servers SET is_latest = false WHERE server_name = $1 AND is_latest = true`
_, err := executor.Exec(ctx, query, serverName)
if err != nil {
return fmt.Errorf("failed to unmark latest version: %w", err)
}
return nil
}
// SetLatestVersion sets is_latest=true on the given version and false on all other versions
// of the same server. Passing an empty version clears is_latest for all rows.
//
// The clear and set are issued as separate statements because the unique partial index
// idx_unique_latest_per_server is non-deferrable and Postgres checks it row-by-row within
// a single UPDATE, which would trip when flipping one row's flag off and another's on.
func (db *PostgreSQL) SetLatestVersion(ctx context.Context, tx pgx.Tx, serverName, version string) error {
if ctx.Err() != nil {
return ctx.Err()
}
executor := db.getExecutor(tx)
if _, err := executor.Exec(ctx,
`UPDATE servers SET is_latest = false WHERE server_name = $1 AND is_latest = true AND version <> $2`,
serverName, version,
); err != nil {
return fmt.Errorf("failed to clear previous latest version: %w", err)
}
if version == "" {
return nil
}
if _, err := executor.Exec(ctx,
`UPDATE servers SET is_latest = true WHERE server_name = $1 AND version = $2 AND is_latest = false`,
serverName, version,
); err != nil {
return fmt.Errorf("failed to set latest version: %w", err)
}
return nil
}
// Close closes the database connection
func (db *PostgreSQL) Close() error {
db.pool.Close()
return nil
}