at-container-registry/pkg/appview/db/go_migrations.go

package db

import (
	"database/sql"
	"fmt"
	"log/slog"
	"strings"
)

// goMigrations holds migration steps that cannot be expressed in SQL, keyed by
// the migration version they belong to.
//
// A hook runs inside the same transaction as its migration's SQL, immediately
// after it and before the version is recorded, so a version is either fully
// applied or not applied at all.
//
// This exists because a migration's data precondition has to be established
// inside the migration sequence. Anything satisfied out-of-band (at runtime, by
// a background worker, by an operator) is invisible to replay: a database
// upgrading across several releases at once runs every pending migration
// back-to-back at boot, long before any worker starts, so a later migration that
// assumed the out-of-band step had happened would silently operate on unprepared
// data.
//
// Keep hooks idempotent. A migration that fails partway is rolled back and
// retried on the next boot, and a hook may also run against a database where
// some or all of its work has already happened by other means.
var goMigrations = map[int]func(*sql.Tx) error{
	33: backfillManifestKeys,
}

// manifestKeyBackfillBatch is how many manifests get their key set per
// statement. Each row contributes three placeholders, so this stays far below
// the driver's parameter ceiling while keeping the number of round trips to a
// remote primary low: one per 500 rows rather than one per row.
const manifestKeyBackfillBatch = 500

// backfillManifestKeys fills manifests.manifest_key for every row that lacks it.
//
// The value is a truncated sha256 of (did, repository, digest). SQLite has no
// hash builtin and go-libsql cannot register one, which is why this is a Go hook
// rather than an UPDATE in 0033's SQL.
//
// On a live database this is usually a no-op or close to it, because the
// Jetstream backfill fills the column as it re-upserts each manifest. It matters
// on the path that has no backfill to rely on: a database replaying migrations
// across releases, where 0033 and anything built on top of it run seconds apart
// with no worker having started. Without this, a later migration that moves
// layers and manifest_references onto manifest_key would copy NULLs.
func backfillManifestKeys(tx *sql.Tx) error {
	type pending struct {
		id  int64
		key string
	}

	var filled int
	for {
		// Read a page, then close the cursor before writing. Updating a table
		// while a SELECT over it is still open is asking for trouble.
		rows, err := tx.Query(`
			SELECT id, did, repository, digest
			FROM manifests
			WHERE manifest_key IS NULL
			LIMIT ?
		`, manifestKeyBackfillBatch)
		if err != nil {
			return fmt.Errorf("select manifests missing a key: %w", err)
		}

		batch := make([]pending, 0, manifestKeyBackfillBatch)
		for rows.Next() {
			var (
				id                      int64
				did, repository, digest string
			)
			if err := rows.Scan(&id, &did, &repository, &digest); err != nil {
				rows.Close()
				return fmt.Errorf("scan manifest: %w", err)
			}
			batch = append(batch, pending{id: id, key: ManifestKey(did, repository, digest)})
		}
		if err := rows.Err(); err != nil {
			rows.Close()
			return fmt.Errorf("iterate manifests: %w", err)
		}
		rows.Close()

		if len(batch) == 0 {
			break
		}

		// One statement per page: UPDATE ... SET manifest_key = CASE id WHEN ...
		// END WHERE id IN (...). A statement per row would be correct too, and
		// unusably slow against a remote primary.
		var sb strings.Builder
		args := make([]any, 0, len(batch)*3)
		sb.WriteString("UPDATE manifests SET manifest_key = CASE id")
		for _, p := range batch {
			sb.WriteString(" WHEN ? THEN ?")
			args = append(args, p.id, p.key)
		}
		sb.WriteString(" END WHERE id IN (")
		for i, p := range batch {
			if i > 0 {
				sb.WriteByte(',')
			}
			sb.WriteByte('?')
			args = append(args, p.id)
		}
		sb.WriteByte(')')

		if _, err := tx.Exec(sb.String(), args...); err != nil {
			return fmt.Errorf("set manifest keys: %w", err)
		}
		filled += len(batch)

		// A short page means the table is exhausted. Checking this rather than
		// looping again saves a query, and the loop still terminates either way
		// because every row written stops matching the WHERE.
		if len(batch) < manifestKeyBackfillBatch {
			break
		}
	}

	if filled > 0 {
		slog.Info("Backfilled manifest_key during migration", "manifests", filled)
	}
	return nil
}