Switch to prefixed IDs and tell agents which Mac a message came from
This commit is contained in:
Ian Gulliver
@@ -0,0 +1,252 @@
|
||||
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, 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, 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
|
||||
}
|
||||
Reference in New Issue
Block a user