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

package db

import (
	"database/sql"
	"os"
	"testing"
	"time"
)

// TestMigrationReplayPopulatesManifestKeys covers the upgrade path that has no
// running application behind it.
//
// A database several releases behind runs every pending migration back-to-back
// at boot, before any worker starts. So a migration cannot depend on a data
// precondition that something else establishes at runtime: from replay's point
// of view that step never happened. manifests.manifest_key is filled by the
// Jetstream backfill on a live system, which is exactly the kind of out-of-band
// step that is invisible here, so 0033 carries a Go hook that fills it in
// sequence instead.
//
// Without the hook, anything built on top of manifest_key (moving layers and
// manifest_references onto it) would silently copy NULLs on this path while
// working fine on a system that had been running.
func TestMigrationReplayPopulatesManifestKeys(t *testing.T) {
	database := baseSchemaDB(t)
	defer database.Close()

	// Seed rows in the pre-0009 shape, as a database at that vintage would have.
	if _, err := database.Exec(
		`INSERT INTO users (did, handle, pds_endpoint, last_seen) VALUES (?, ?, ?, ?)`,
		"did:plc:replay", "replay.example.com", "https://pds.example.com", time.Now(),
	); err != nil {
		t.Fatalf("seed user: %v", err)
	}

	seeded := []struct{ repo, digest string }{
		{"myapp", "sha256:aaa"},
		{"myapp", "sha256:bbb"},
		{"otherapp", "sha256:ccc"},
	}
	for _, s := range seeded {
		if _, err := database.Exec(`
			INSERT INTO manifests
			(did, repository, digest, hold_endpoint, schema_version, media_type, artifact_type, created_at)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?)
		`, "did:plc:replay", s.repo, s.digest, "did:web:hold.example.com", 2,
			"application/vnd.oci.image.manifest.v1+json", "container-image", time.Now()); err != nil {
			t.Fatalf("seed manifest %s/%s: %v", s.repo, s.digest, err)
		}
	}

	// Replay everything from that vintage forward, with no application running.
	if err := runMigrations(database, false); err != nil {
		t.Fatalf("replay migrations: %v", err)
	}

	var total int
	if err := database.QueryRow(`SELECT COUNT(*) FROM manifests`).Scan(&total); err != nil {
		t.Fatalf("count manifests: %v", err)
	}
	if total != len(seeded) {
		t.Fatalf("manifest count = %d after replay, want %d", total, len(seeded))
	}

	for _, s := range seeded {
		var key string
		if err := database.QueryRow(
			`SELECT manifest_key FROM manifests WHERE did = ? AND repository = ? AND digest = ?`,
			"did:plc:replay", s.repo, s.digest,
		).Scan(&key); err != nil {
			t.Fatalf("read key for %s/%s: %v", s.repo, s.digest, err)
		}
		if want := ManifestKey("did:plc:replay", s.repo, s.digest); key != want {
			t.Errorf("%s/%s key = %q, want %q", s.repo, s.digest, key, want)
		}
	}
}

// TestManifestKeyHookIsIdempotent: a migration that fails partway is rolled back
// and retried on the next boot, and the hook also runs against databases where
// the runtime backfill already did some of the work. Neither may break it.
func TestManifestKeyHookIsIdempotent(t *testing.T) {
	database := revTestDB(t)
	reshapeManifestsToPre0034(t, database)
	seedPre0034Manifests(t, database, "did:plc:alice", []string{"sha256:aaa", "sha256:bbb"})

	// One row already filled, one not: the mixed state a partly-backfilled
	// database is actually in.
	if _, err := database.Exec(
		`UPDATE manifests SET manifest_key = ? WHERE digest = ?`,
		ManifestKey("did:plc:alice", "myapp", "sha256:bbb"), "sha256:bbb"); err != nil {
		t.Fatalf("pre-fill one key: %v", err)
	}

	for pass := range 2 {
		tx, err := database.Begin()
		if err != nil {
			t.Fatalf("begin: %v", err)
		}
		if err := backfillManifestKeys(tx); err != nil {
			tx.Rollback()
			t.Fatalf("pass %d: %v", pass, err)
		}
		if err := tx.Commit(); err != nil {
			t.Fatalf("commit pass %d: %v", pass, err)
		}
	}

	if n := countNullKeys(t, database); n != 0 {
		t.Errorf("%d manifests still unfilled after two passes", n)
	}
	for _, digest := range []string{"sha256:aaa", "sha256:bbb"} {
		var got string
		if err := database.QueryRow(`SELECT manifest_key FROM manifests WHERE digest = ?`, digest).Scan(&got); err != nil {
			t.Fatalf("read %s: %v", digest, err)
		}
		if want := ManifestKey("did:plc:alice", "myapp", digest); got != want {
			t.Errorf("%s key = %q, want %q", digest, got, want)
		}
	}
}

// TestManifestKeyHookSpansPages: the hook pages through the table, so the
// interesting case is more rows than one page.
func TestManifestKeyHookSpansPages(t *testing.T) {
	database := revTestDB(t)
	reshapeManifestsToPre0034(t, database)

	const total = manifestKeyBackfillBatch + 21
	digests := make([]string, total)
	for i := range total {
		digests[i] = digestFor(i)
	}
	seedPre0034Manifests(t, database, "did:plc:alice", digests)

	tx, err := database.Begin()
	if err != nil {
		t.Fatalf("begin: %v", err)
	}
	if err := backfillManifestKeys(tx); err != nil {
		tx.Rollback()
		t.Fatalf("backfillManifestKeys: %v", err)
	}
	if err := tx.Commit(); err != nil {
		t.Fatalf("commit: %v", err)
	}

	if n := countNullKeys(t, database); n != 0 {
		t.Errorf("%d of %d unfilled; the hook stopped before the last page", n, total)
	}
}

func digestFor(i int) string {
	const hexdigits = "0123456789abcdef"
	out := []byte("sha256:0000000000000000")
	for pos, n := len(out)-1, i; pos >= 7 && n > 0; pos, n = pos-1, n/16 {
		out[pos] = hexdigits[n%16]
	}
	return string(out)
}

// baseSchemaDB returns a database at the pre-0009 vintage, with no migrations
// applied beyond the ones that snapshot records.
func baseSchemaDB(t *testing.T) *sql.DB {
	t.Helper()

	base, err := os.ReadFile("testdata/base_schema.sql")
	if err != nil {
		t.Fatalf("read testdata/base_schema.sql: %v", err)
	}
	database := revTestDBRaw(t)
	for i, stmt := range splitSQLStatements(string(base)) {
		if _, err := database.Exec(stmt); err != nil {
			t.Fatalf("apply base_schema.sql statement %d: %v", i+1, err)
		}
	}
	return database
}

// revTestDBRaw opens an empty file-backed database with no schema applied, so
// the caller can build one from a snapshot.
func revTestDBRaw(t *testing.T) *sql.DB {
	t.Helper()
	connector, err := openLibsqlLocalConnector("file:" + t.TempDir() + "/replay.db")
	if err != nil {
		t.Fatalf("open connector: %v", err)
	}
	database := sql.OpenDB(connector)
	database.SetMaxOpenConns(1)
	t.Cleanup(func() { database.Close() })
	return database
}

// countNullKeys reports how many manifests still lack a key. Only meaningful on
// a database that has not yet reached migration 0034, where manifest_key becomes
// the primary key and can no longer be null.
func countNullKeys(t *testing.T, database *sql.DB) int {
	t.Helper()
	var n int
	if err := database.QueryRow(`SELECT COUNT(*) FROM manifests WHERE manifest_key IS NULL`).Scan(&n); err != nil {
		t.Fatalf("count unfilled keys: %v", err)
	}
	return n
}

// reshapeManifestsToPre0034 rebuilds the manifests table in its 0033 shape: a
// surrogate id, with manifest_key present but nullable. backfillManifestKeys only
// ever runs at that point in the migration sequence, so testing it against the
// current schema would be testing something that cannot happen.
func reshapeManifestsToPre0034(t *testing.T, database *sql.DB) {
	t.Helper()
	stmts := []string{
		`DROP TABLE layers`,
		`DROP TABLE manifest_references`,
		`DROP TABLE manifests`,
		`CREATE TABLE manifests (
			id INTEGER PRIMARY KEY AUTOINCREMENT,
			did TEXT NOT NULL,
			repository TEXT NOT NULL,
			digest TEXT NOT NULL,
			hold_endpoint TEXT NOT NULL,
			schema_version INTEGER NOT NULL,
			media_type TEXT NOT NULL,
			config_digest TEXT,
			config_size INTEGER,
			artifact_type TEXT NOT NULL DEFAULT 'container-image',
			subject_digest TEXT,
			created_at TIMESTAMP NOT NULL,
			manifest_key TEXT,
			UNIQUE(did, repository, digest)
		)`,
	}
	for _, stmt := range stmts {
		if _, err := database.Exec(stmt); err != nil {
			t.Fatalf("reshape manifests: %v", err)
		}
	}
}

// seedPre0034Manifests inserts manifests with no key, as rows predating 0033
// would be.
func seedPre0034Manifests(t *testing.T, database *sql.DB, did string, digests []string) {
	t.Helper()
	for _, digest := range digests {
		if _, err := database.Exec(`
			INSERT INTO manifests (did, repository, digest, hold_endpoint, schema_version, media_type, artifact_type, created_at)
			VALUES (?, 'myapp', ?, 'did:web:hold', 2, 'application/vnd.oci.image.manifest.v1+json', 'container-image', ?)
		`, did, digest, time.Now()); err != nil {
			t.Fatalf("seed %s: %v", digest, err)
		}
	}
}