Something went wrong. Try again.
Monorepo for Tangled
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209package db
import ( "database/sql" "encoding/json" "fmt" "slices" "strings" "time")
// the raw token is never stored, only its hashtype ExecutorRegistration struct { Name string CreatedAt string ExpiresAt *time.Time Labels []string}
func (d *DB) CreateExecutorToken(tokenHash string, expiresAt *time.Time) error { _, err := d.Exec( `insert into mill_executor_tokens (token_hash, expires_at) values (?, ?)`, tokenHash, expiryArg(expiresAt), ) return err}
func (d *DB) RegisterExecutor(name, tokenHash string, labels []string) error { normalized := normalizeLabels(labels) labelsBytes, err := json.Marshal(normalized) if err != nil { return err }
tx, err := d.Begin() if err != nil { return err } defer tx.Rollback()
var expires sql.NullString err = tx.QueryRow( `select expires_at from mill_executor_tokens where token_hash = ?`, tokenHash, ).Scan(&expires) if err == sql.ErrNoRows { return fmt.Errorf("unknown executor token") } if err != nil { return err } exp, hasExpiry, err := parseExpiry(expires) if err != nil { return err } if hasExpiry && time.Now().After(exp) { return fmt.Errorf("executor token expired at %s", exp.Format(time.RFC3339)) }
var registeredName string err = tx.QueryRow( `select name from mill_executors where token_hash = ?`, tokenHash, ).Scan(®isteredName) if err != nil && err != sql.ErrNoRows { return err } if err == nil && registeredName != name { return fmt.Errorf("executor token is already registered to %q", registeredName) }
res, err := tx.Exec(` insert into mill_executors (name, token_hash, labels) values (?, ?, ?) on conflict(name) do update set labels = excluded.labels where mill_executors.token_hash = excluded.token_hash `, name, tokenHash, string(labelsBytes)) if err != nil { return err } changed, err := res.RowsAffected() if err != nil { return err } if changed == 0 { return fmt.Errorf("executor %q is already registered with another token", name) } return tx.Commit()}
func (d *DB) ResolveExecutorToken(tokenHash string) (string, []string, bool, error) { var name string var expires, labelsRaw sql.NullString err := d.QueryRow(` select e.name, t.expires_at, e.labels from mill_executors e join mill_executor_tokens t on t.token_hash = e.token_hash where e.token_hash = ? `, tokenHash).Scan(&name, &expires, &labelsRaw) if err == sql.ErrNoRows { return "", nil, false, nil } if err != nil { return "", nil, false, err } exp, hasExpiry, err := parseExpiry(expires) if err != nil { return "", nil, false, err } if hasExpiry && time.Now().After(exp) { return "", nil, false, nil } var labels []string if labelsRaw.Valid && labelsRaw.String != "" { if err := json.Unmarshal([]byte(labelsRaw.String), &labels); err != nil { return "", nil, false, err } } return name, labels, true, nil}
func (d *DB) RevokeExecutorToken(name string) (bool, error) { res, err := d.Exec(` delete from mill_executor_tokens where token_hash = (select token_hash from mill_executors where name = ?) `, name) if err != nil { return false, err } n, err := res.RowsAffected() return n > 0, err}
func (d *DB) ListExecutorRegistrations() ([]ExecutorRegistration, error) { rows, err := d.Query(` select e.name, t.created_at, t.expires_at, e.labels from mill_executors e join mill_executor_tokens t on t.token_hash = e.token_hash order by e.name `) if err != nil { return nil, err } defer rows.Close()
var out []ExecutorRegistration for rows.Next() { var registration ExecutorRegistration var expires, labelsRaw sql.NullString if err := rows.Scan( ®istration.Name, ®istration.CreatedAt, &expires, &labelsRaw, ); err != nil { return nil, err } exp, hasExpiry, err := parseExpiry(expires) if err != nil { return nil, err } if hasExpiry { registration.ExpiresAt = &exp } if labelsRaw.Valid && labelsRaw.String != "" { if err := json.Unmarshal([]byte(labelsRaw.String), ®istration.Labels); err != nil { return nil, err } } out = append(out, registration) } return out, rows.Err()}
func normalizeLabels(labels []string) []string { var out []string seen := make(map[string]bool) for _, l := range labels { trimmed := strings.TrimSpace(l) if trimmed == "" { continue } if !seen[trimmed] { seen[trimmed] = true out = append(out, trimmed) } } slices.Sort(out) return out}
func expiryArg(t *time.Time) any { if t == nil { return nil } return t.UTC().Format(time.RFC3339)}
func parseExpiry(s sql.NullString) (time.Time, bool, error) { if !s.Valid { return time.Time{}, false, nil } t, err := time.Parse(time.RFC3339, s.String) if err != nil { return time.Time{}, false, err } return t, true, nil}