Add multi-account usage dashboard
This commit is contained in:
@@ -42,7 +42,138 @@ CREATE TABLE IF NOT EXISTS notifications (dedupe_key TEXT PRIMARY KEY, channel T
|
||||
CREATE TABLE IF NOT EXISTS telegram_updates (id INTEGER PRIMARY KEY CHECK(id=1), offset INTEGER NOT NULL DEFAULT 0);
|
||||
INSERT OR IGNORE INTO telegram_updates(id,offset) VALUES(1,0);
|
||||
`)
|
||||
return err
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.migrateAccounts()
|
||||
}
|
||||
|
||||
func (s *Store) migrateAccounts() error {
|
||||
tx, err := s.DB.Begin()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
if _, err = tx.Exec(`CREATE TABLE IF NOT EXISTS accounts (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
display_name TEXT NOT NULL,
|
||||
email TEXT,
|
||||
plan_type TEXT,
|
||||
connected INTEGER NOT NULL DEFAULT 0,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL
|
||||
)`); err != nil {
|
||||
return err
|
||||
}
|
||||
var count int
|
||||
if err = tx.QueryRow("SELECT COUNT(*) FROM accounts").Scan(&count); err != nil {
|
||||
return err
|
||||
}
|
||||
if count == 0 {
|
||||
if _, err = tx.Exec("INSERT INTO accounts(id,display_name,created_at,updated_at) VALUES(1,'默认账号',?,?)", time.Now().Unix(), time.Now().Unix()); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
for _, table := range []string{"daily_usage", "limit_snapshots"} {
|
||||
var found int
|
||||
rows, qerr := tx.Query("PRAGMA table_info(" + table + ")")
|
||||
if qerr != nil {
|
||||
return qerr
|
||||
}
|
||||
for rows.Next() {
|
||||
var cid, notnull, pk int
|
||||
var name, typ string
|
||||
var def any
|
||||
_ = rows.Scan(&cid, &name, &typ, ¬null, &def, &pk)
|
||||
if name == "account_id" {
|
||||
found = 1
|
||||
}
|
||||
}
|
||||
rows.Close()
|
||||
if found == 0 {
|
||||
if table == "daily_usage" {
|
||||
_, err = tx.Exec(`ALTER TABLE daily_usage RENAME TO daily_usage_legacy;
|
||||
CREATE TABLE daily_usage (account_id INTEGER NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,date TEXT NOT NULL,total_tokens INTEGER NOT NULL,fetched_at INTEGER NOT NULL,PRIMARY KEY(account_id,date));
|
||||
INSERT INTO daily_usage SELECT 1,date,total_tokens,fetched_at FROM daily_usage_legacy;
|
||||
DROP TABLE daily_usage_legacy;`)
|
||||
} else {
|
||||
_, err = tx.Exec(`ALTER TABLE limit_snapshots RENAME TO limit_snapshots_legacy;
|
||||
CREATE TABLE limit_snapshots (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
account_id INTEGER NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
|
||||
limit_id TEXT NOT NULL,
|
||||
window_type TEXT NOT NULL,
|
||||
used_percent REAL NOT NULL,
|
||||
duration_mins INTEGER NOT NULL,
|
||||
resets_at INTEGER NOT NULL,
|
||||
fetched_at INTEGER NOT NULL
|
||||
);
|
||||
INSERT INTO limit_snapshots(id,account_id,limit_id,window_type,used_percent,duration_mins,resets_at,fetched_at)
|
||||
SELECT id,1,limit_id,window_type,used_percent,duration_mins,resets_at,fetched_at FROM limit_snapshots_legacy;
|
||||
DROP TABLE limit_snapshots_legacy;
|
||||
CREATE INDEX idx_limits_time ON limit_snapshots(fetched_at);`)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
}
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
type Account struct {
|
||||
ID int64 `json:"id"`
|
||||
DisplayName string `json:"displayName"`
|
||||
Email *string `json:"email"`
|
||||
PlanType *string `json:"planType"`
|
||||
Connected bool `json:"connected"`
|
||||
CreatedAt int64 `json:"createdAt"`
|
||||
UpdatedAt int64 `json:"updatedAt"`
|
||||
}
|
||||
|
||||
func (s *Store) Accounts() ([]Account, error) {
|
||||
rows, e := s.DB.Query("SELECT id,display_name,email,plan_type,connected,created_at,updated_at FROM accounts ORDER BY id")
|
||||
if e != nil {
|
||||
return nil, e
|
||||
}
|
||||
defer rows.Close()
|
||||
out := []Account{}
|
||||
for rows.Next() {
|
||||
var a Account
|
||||
if e = rows.Scan(&a.ID, &a.DisplayName, &a.Email, &a.PlanType, &a.Connected, &a.CreatedAt, &a.UpdatedAt); e != nil {
|
||||
return nil, e
|
||||
}
|
||||
out = append(out, a)
|
||||
}
|
||||
return out, rows.Err()
|
||||
}
|
||||
func (s *Store) CreateAccount(name string) (Account, error) {
|
||||
now := time.Now().Unix()
|
||||
r, e := s.DB.Exec("INSERT INTO accounts(display_name,created_at,updated_at) VALUES(?,?,?)", name, now, now)
|
||||
if e != nil {
|
||||
return Account{}, e
|
||||
}
|
||||
id, _ := r.LastInsertId()
|
||||
return Account{ID: id, DisplayName: name, CreatedAt: now, UpdatedAt: now}, nil
|
||||
}
|
||||
func (s *Store) RenameAccount(id int64, name string) error {
|
||||
r, e := s.DB.Exec("UPDATE accounts SET display_name=?,updated_at=? WHERE id=?", name, time.Now().Unix(), id)
|
||||
if e != nil {
|
||||
return e
|
||||
}
|
||||
n, _ := r.RowsAffected()
|
||||
if n == 0 {
|
||||
return sql.ErrNoRows
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (s *Store) UpdateAccount(id int64, email, plan *string, connected bool) error {
|
||||
_, e := s.DB.Exec("UPDATE accounts SET email=?,plan_type=?,connected=?,updated_at=? WHERE id=?", email, plan, connected, time.Now().Unix(), id)
|
||||
return e
|
||||
}
|
||||
func (s *Store) DeleteAccount(id int64) error {
|
||||
_, e := s.DB.Exec("DELETE FROM accounts WHERE id=?", id)
|
||||
return e
|
||||
}
|
||||
|
||||
func (s *Store) Get(key string) (string, bool) {
|
||||
|
||||
@@ -26,6 +26,104 @@ func TestSettings(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountsAndPerAccountUsage(t *testing.T) {
|
||||
s, err := Open(t.TempDir())
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.DB.Close()
|
||||
accounts, err := s.Accounts()
|
||||
if err != nil || len(accounts) != 1 || accounts[0].ID != 1 {
|
||||
t.Fatalf("default accounts = %#v, %v", accounts, err)
|
||||
}
|
||||
second, err := s.CreateAccount("Team workspace")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, id := range []int64{1, second.ID} {
|
||||
if _, err = s.DB.Exec("INSERT INTO daily_usage(account_id,date,total_tokens,fetched_at) VALUES(?,?,?,?)", id, "2026-08-13", id*100, 1); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if err = s.DeleteAccount(second.ID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var count int
|
||||
if err = s.DB.QueryRow("SELECT COUNT(*) FROM daily_usage WHERE account_id=?", second.ID).Scan(&count); err != nil || count != 0 {
|
||||
t.Fatalf("usage was not cascaded: %d, %v", count, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLegacyUsageMigratesToDefaultAccount(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
db, err := sql.Open("sqlite", filepath.Join(dir, "codex-helper.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err = db.Exec("CREATE TABLE daily_usage(date TEXT PRIMARY KEY,total_tokens INTEGER NOT NULL,fetched_at INTEGER NOT NULL); INSERT INTO daily_usage VALUES('2026-08-12',321,1)"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db.Close()
|
||||
s, err := Open(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.DB.Close()
|
||||
var accountID, tokens int64
|
||||
if err = s.DB.QueryRow("SELECT account_id,total_tokens FROM daily_usage WHERE date='2026-08-12'").Scan(&accountID, &tokens); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if accountID != 1 || tokens != 321 {
|
||||
t.Fatalf("migrated row = account %d, tokens %d", accountID, tokens)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPopulatedLegacyLimitSnapshotsMigrateToDefaultAccount(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
db, err := sql.Open("sqlite", filepath.Join(dir, "codex-helper.db"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
_, err = db.Exec(`
|
||||
CREATE TABLE limit_snapshots (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
limit_id TEXT NOT NULL,
|
||||
window_type TEXT NOT NULL,
|
||||
used_percent REAL NOT NULL,
|
||||
duration_mins INTEGER NOT NULL,
|
||||
resets_at INTEGER NOT NULL,
|
||||
fetched_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX idx_limits_time ON limit_snapshots(fetched_at);
|
||||
INSERT INTO limit_snapshots(id,limit_id,window_type,used_percent,duration_mins,resets_at,fetched_at)
|
||||
VALUES(7,'codex','primary',42.5,300,1700000000,1699990000);
|
||||
`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db.Close()
|
||||
|
||||
s, err := Open(dir)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer s.DB.Close()
|
||||
var id, accountID, duration int64
|
||||
var limitID, window string
|
||||
var used float64
|
||||
err = s.DB.QueryRow("SELECT id,account_id,limit_id,window_type,used_percent,duration_mins FROM limit_snapshots").Scan(&id, &accountID, &limitID, &window, &used, &duration)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if id != 7 || accountID != 1 || limitID != "codex" || window != "primary" || used != 42.5 || duration != 300 {
|
||||
t.Fatalf("migrated limit = id %d, account %d, %s/%s, %.1f, duration %d", id, accountID, limitID, window, used, duration)
|
||||
}
|
||||
var indexCount int
|
||||
if err = s.DB.QueryRow("SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name='idx_limits_time'").Scan(&indexCount); err != nil || indexCount != 1 {
|
||||
t.Fatalf("limit index was not recreated: %d, %v", indexCount, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBackupIncludesCommittedWALData(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s, err := Open(dir)
|
||||
|
||||
Reference in New Issue
Block a user