byre/internal/store/migrateids.go

package store

import (
	"database/sql"
	"errors"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"g.fc.run/ai/byre/internal/git"
	"g.fc.run/ai/byre/internal/ids"
)

var idToken = regexp.MustCompile(`[0-9A-Fa-f]{8}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{4}-[0-9A-Fa-f]{12}|queued-[0-9]+|[0-9A-Za-z]+`)

var keptColumns = map[string]bool{
	"agent_sessions.vendor_session_id": true,
	"agent_sessions.forked_from":       true,
}

func migratePrefixIDs(tx *sql.Tx, dataDir string) error {
	renames, err := idRenames(tx)
	if err != nil {
		return err
	}
	rewrite := func(s string) string {
		return idToken.ReplaceAllStringFunc(s, func(tok string) string {
			if to, ok := renames[tok]; ok {
				return to
			}
			return tok
		})
	}

	tables, err := textColumns(tx)
	if err != nil {
		return err
	}
	for table, columns := range tables {
		err = rewriteTable(tx, table, columns, rewrite)
		if err != nil {
			return fmt.Errorf("%s: %w", table, err)
		}
	}

	repos := map[string]string{}
	rows, err := tx.Query(`SELECT id, repo_path FROM projects`)
	if err != nil {
		return err
	}
	for rows.Next() {
		id, repo := "", ""
		err = rows.Scan(&id, &repo)
		if err != nil {
			rows.Close()
			return err
		}
		repos[id] = repo
	}
	rows.Close()
	err = rows.Err()
	if err != nil {
		return err
	}

	for _, sub := range []string{"blobs", "hands", "antigravity", "worktrees"} {
		err = renameEntries(filepath.Join(dataDir, sub), rewrite)
		if err != nil {
			return err
		}
	}
	for project, repo := range repos {
		dir := filepath.Join(dataDir, "worktrees", project)
		err = renameEntries(dir, rewrite)
		if err != nil {
			return err
		}
		worktrees, err := filepath.Glob(filepath.Join(dir, "*", ".git"))
		if err != nil {
			return err
		}
		for i, w := range worktrees {
			worktrees[i] = filepath.Dir(w)
		}
		err = git.RepairWorktrees(repo, worktrees)
		if err != nil {
			return err
		}
		err = git.RenameRefs(repo, rewrite)
		if err != nil {
			return err
		}
		err = git.RewriteNotes(repo, rewrite)
		if err != nil {
			return err
		}
	}
	return nil
}

func idRenames(tx *sql.Tx) (map[string]string, error) {
	sources := []struct {
		prefix ids.Prefix
		query  string
	}{
		{ids.Project, `SELECT id FROM projects`},
		{ids.Task, `SELECT id FROM tasks`},
		{ids.WorkflowRun, `SELECT DISTINCT workflow_run FROM tasks`},
		{ids.Session, `SELECT id FROM agent_sessions`},
		{ids.Snapshot, `SELECT id FROM snapshots`},
		{ids.Review, `SELECT id FROM reviews`},
		{ids.InboxItem, `SELECT id FROM inbox_items`},
		{ids.InboxResponse, `SELECT id FROM inbox_responses`},
		{ids.Blob, `SELECT id FROM blobs`},
		{ids.LeadNote, `SELECT id FROM lead_notes`},
		{ids.WorkflowDefinition, `SELECT id FROM workflow_definitions`},
		{ids.Send, `SELECT DISTINCT send_id FROM task_inputs`},
		{ids.Send, `SELECT DISTINCT send_id FROM lead_inputs`},
		{ids.Send, `SELECT DISTINCT json_extract(payload, '$.send_id') FROM events WHERE kind IN ('user.queued', 'user.message') AND json_extract(payload, '$.send_id') IS NOT NULL`},
	}

	renames := map[string]string{}
	for _, src := range sources {
		rows, err := tx.Query(src.query)
		if err != nil {
			return nil, err
		}
		for rows.Next() {
			old := ""
			err = rows.Scan(&old)
			if err != nil {
				rows.Close()
				return nil, err
			}
			_, mapped := renames[old]
			if old == "" || mapped || strings.HasPrefix(old, string(src.prefix)) {
				continue
			}
			renames[old] = ids.Derive(src.prefix, old)
		}
		rows.Close()
		err = rows.Err()
		if err != nil {
			return nil, err
		}
	}
	return renames, nil
}

func textColumns(tx *sql.Tx) (map[string][]string, error) {
	rows, err := tx.Query(`SELECT m.name, p.name FROM sqlite_master m JOIN pragma_table_info(m.name) p WHERE m.type = 'table' AND m.name NOT IN ('migrations', 'sqlite_sequence') AND upper(p.type) = 'TEXT'`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	tables := map[string][]string{}
	for rows.Next() {
		table, column := "", ""
		err = rows.Scan(&table, &column)
		if err != nil {
			return nil, err
		}
		if keptColumns[table+"."+column] {
			continue
		}
		tables[table] = append(tables[table], column)
	}
	return tables, rows.Err()
}

func rewriteTable(tx *sql.Tx, table string, columns []string, rewrite func(string) string) error {
	quoted := []string{}
	for _, c := range columns {
		quoted = append(quoted, `"`+c+`"`)
	}
	query := fmt.Sprintf(`SELECT rowid, %s FROM "%s" WHERE rowid > ? ORDER BY rowid LIMIT 2000`, strings.Join(quoted, ", "), table)

	after := int64(0)
	for {
		type change struct {
			rowid  int64
			column string
			value  string
		}
		changes := []change{}
		n := 0

		rows, err := tx.Query(query, after)
		if err != nil {
			return err
		}
		for rows.Next() {
			values := make([]sql.NullString, len(columns))
			dest := []any{&after}
			for i := range values {
				dest = append(dest, &values[i])
			}
			err = rows.Scan(dest...)
			if err != nil {
				rows.Close()
				return err
			}
			n++
			for i, v := range values {
				if !v.Valid {
					continue
				}
				to := rewrite(v.String)
				if to != v.String {
					changes = append(changes, change{after, columns[i], to})
				}
			}
		}
		rows.Close()
		err = rows.Err()
		if err != nil {
			return err
		}

		for _, c := range changes {
			_, err = tx.Exec(fmt.Sprintf(`UPDATE "%s" SET "%s" = ? WHERE rowid = ?`, table, c.column), c.value, c.rowid)
			if err != nil {
				return err
			}
		}
		if n == 0 {
			return nil
		}
	}
}

func renameEntries(dir string, rewrite func(string) string) error {
	entries, err := os.ReadDir(dir)
	if errors.Is(err, fs.ErrNotExist) {
		return nil
	}
	if err != nil {
		return err
	}
	for _, e := range entries {
		to := rewrite(e.Name())
		if to == e.Name() {
			continue
		}
		err = os.Rename(filepath.Join(dir, e.Name()), filepath.Join(dir, to))
		if err != nil {
			return err
		}
	}
	return nil
}