llmgw/internal/store/copy.go

package store

import (
	"context"
	"fmt"
	"strings"
)

// copyOrder is every table, parents before children.
//
// Written out rather than discovered, because the order is a fact about the
// foreign keys and getting it wrong produces a constraint violation halfway
// through — with some tables copied and some not, which is the worst place to
// stop. A table added without being added here is copied by nothing, and the
// destination comes up quietly missing it; the test below is what catches
// that.
var copyOrder = []string{
	"organizations",
	"teams",
	"projects",
	// Before api_keys, which references it. The first version of this list
	// followed the order the tables are created across the migrations, and
	// users is created in a later one than api_keys even though api_keys
	// points at it — so the copy failed on a foreign key three tables in.
	"users",
	// After users and the three levels of the tree, which it points at.
	"memberships",
	"api_keys",
	"key_claims",
	"sessions",
	"limits",
	"usage",
	"models",
	// No foreign key to models, deliberately: a measurement outlives the
	// catalog entry it was taken for, and losing it when somebody removes a
	// model would mean re-running a benchmark that costs real GPU time.
	"capacity",
	"threads",
	"thread_messages",
	"memory_notes",
}

// notCopied are tables deliberately left behind.
//
// Leadership is a claim about a process that is running right now. Copying it
// would hand the new store a leader that does not exist, and the standby
// waiting on that lease would wait out its full expiry before noticing — at
// the exact moment somebody is watching a migration and wondering why nothing
// is being placed.
var notCopied = map[string]string{
	"leadership": "a lease on a running process, meaningless in a new store",
}

// CopyResult is how many rows each table received.
type CopyResult struct {
	Table string
	Rows  int64
}

// CopyTo copies every row into another store.
//
// For moving from SQLite to PostgreSQL, which otherwise means reissuing every
// API key: the keys are stored hashed, so there is no way to recreate them and
// no way for anyone to keep using the one they have.
//
// The destination must be empty. Merging two stores would mean deciding what
// happens when both have a project with the same id and different contents,
// and every answer to that is a guess about which one someone meant.
//
// Not a transaction across the whole copy. It is one transaction per table,
// because a single transaction holding every row of a large usage table is how
// a migration runs out of memory on the machine least able to spare it. A
// failure therefore leaves a partial copy, which is why the destination has to
// have been empty: starting again is dropping and recreating it.
func (s *DB) CopyTo(ctx context.Context, dst *DB) ([]CopyResult, error) {
	empty, where, err := dst.isEmpty(ctx)
	if err != nil {
		return nil, err
	}
	if !empty {
		return nil, fmt.Errorf("the destination already has rows in %s; copy into an empty store", where)
	}

	var out []CopyResult
	for _, table := range copyOrder {
		n, err := s.copyTable(ctx, dst, table)
		if err != nil {
			return out, fmt.Errorf("copying %s: %w", table, err)
		}
		out = append(out, CopyResult{Table: table, Rows: n})
	}
	// Part of the copy, not a step someone has to remember: a destination
	// whose sequences still start at 1 is not a destination anyone can use.
	if err := dst.resetIdentities(ctx); err != nil {
		return out, fmt.Errorf("pointing the identity columns past the copied rows: %w", err)
	}
	return out, nil
}

// isEmpty reports whether any table has rows, and which one first does.
// Empty reports whether the store holds no rows at all, which is what CopyTo
// needs of its destination.
func (s *DB) Empty(ctx context.Context) (bool, error) {
	empty, _, err := s.isEmpty(ctx)
	return empty, err
}

func (s *DB) isEmpty(ctx context.Context) (bool, string, error) {
	for _, table := range copyOrder {
		var n int64
		//nolint:gosec // table comes from copyOrder, not from input
		if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM `+table).Scan(&n); err != nil {
			return false, "", err
		}
		if n > 0 {
			return false, table, nil
		}
	}
	return true, "", nil
}

// copyTable moves one table, column names taken from the source.
//
// Reading the columns rather than listing them keeps this working when a
// migration adds one: a hand-written column list would copy the old shape and
// silently drop the new field, and nothing would notice until someone looked
// for a setting that was there before the move.
func (s *DB) copyTable(ctx context.Context, dst *DB, table string) (int64, error) {
	//nolint:gosec // table comes from copyOrder, not from input
	rows, err := s.db.QueryContext(ctx, `SELECT * FROM `+table)
	if err != nil {
		return 0, err
	}
	defer rows.Close()

	cols, err := rows.Columns()
	if err != nil {
		return 0, err
	}

	q, tx, err := dst.begin(ctx)
	if err != nil {
		return 0, err
	}
	defer func() { _ = tx.Rollback() }()

	insert := `INSERT INTO ` + table + ` (` + strings.Join(quoteAll(cols), ", ") +
		`) VALUES (` + strings.Repeat("?, ", len(cols)-1) + `?)`

	var n int64
	for rows.Next() {
		vals := make([]any, len(cols))
		ptrs := make([]any, len(cols))
		for i := range vals {
			ptrs[i] = &vals[i]
		}
		if err := rows.Scan(ptrs...); err != nil {
			return n, err
		}
		for i, v := range vals {
			vals[i] = keepEmptyBytes(v)
		}
		if _, err := q.ExecContext(ctx, insert, vals...); err != nil {
			return n, err
		}
		n++
	}
	if err := rows.Err(); err != nil {
		return n, err
	}
	return n, tx.Commit()
}

// keepEmptyBytes stops an empty byte column becoming NULL on the way across.
//
// An empty BLOB scans as a typed nil slice — []byte(nil) inside a non-nil any
// — while a real NULL scans as an untyped nil. The two are distinguishable
// here and not in pgx, which encodes any nil slice as NULL, so the difference
// has to be restored before the insert.
//
// It is not a detail. key_claims.sealed is emptied the moment a claim is used,
// so every claimed row holds an empty blob in a NOT NULL column, and the copy
// failed on the first one — after the gateway had been stopped and four tables
// had already been written.
func keepEmptyBytes(v any) any {
	if b, ok := v.([]byte); ok && b == nil {
		return []byte{}
	}
	return v
}

// quoteAll quotes column names, because `key` and `text` are reserved words in
// one dialect or the other and both are column names here.
func quoteAll(cols []string) []string {
	out := make([]string, len(cols))
	for i, c := range cols {
		out[i] = `"` + c + `"`
	}
	return out
}

// resetIdentities points PostgreSQL's identity columns past the copied rows.
//
// Without it the first insert after a migration collides with a row that was
// copied in: the sequence still starts at 1 while the table already holds a
// thousand notes, and the failure arrives later, to somebody else, looking
// like a bug in the feature rather than in the move.
func (s *DB) resetIdentities(ctx context.Context) error {
	if s.dialect != PostgresDialect {
		return nil
	}
	// Only the tables with a generated id. SQLite's rowid needs nothing.
	for _, t := range []string{"memory_notes", "usage"} {
		var has bool
		if err := s.db.QueryRowContext(ctx,
			`SELECT EXISTS (SELECT 1 FROM information_schema.columns
			                WHERE table_name = ? AND column_name = 'id'
			                  AND is_identity = 'YES')`, t).Scan(&has); err != nil {
			return err
		}
		if !has {
			continue
		}
		//nolint:gosec // t is a literal from the list above
		if _, err := s.db.ExecContext(ctx,
			`SELECT setval(pg_get_serial_sequence('`+t+`', 'id'),
			               GREATEST((SELECT COALESCE(MAX(id), 0) FROM `+t+`), 1))`); err != nil {
			return err
		}
	}
	return nil
}