m1saka
fa9640d919
build / build (push) Successful in 2m34s
- New internal/stats (types) and internal/store (SQLite owner: requests + credentials tables, WAL); store implements stats.Recorder. - Stream (SSE tee) and non-stream chat paths parse upstream usage incl. cached_tokens and record per-request; add /api/stats, /api/stats/reset, /admin/stats HTML with cache hit rate. - Drop credentials.json: remove auth file I/O and ZHANLU_CREDENTIALS_FILE; credential precedence is env vars > db row.
90 lines
3.0 KiB
Go
90 lines
3.0 KiB
Go
// Package store owns the single embedded SQLite database backing the zhanlu
|
|
// proxy: token-usage records (the requests table) and the persisted login
|
|
// credentials (the credentials table, single-tenant single row). It implements
|
|
// stats.Recorder for usage tracking and exposes Load/Save for credentials.
|
|
package store
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"strings"
|
|
|
|
_ "modernc.org/sqlite"
|
|
)
|
|
|
|
// Store is the single owner of the proxy's SQLite database handle.
|
|
type Store struct {
|
|
db *sql.DB
|
|
}
|
|
|
|
// Open opens (or creates) the database at path and ensures both tables exist.
|
|
// SQLite is opened with WAL journaling and a busy timeout so concurrent reads
|
|
// (stats queries) and writes (request records, credential saves) do not
|
|
// collide.
|
|
func Open(path string) (*Store, error) {
|
|
db, err := sql.Open("sqlite", path+"?_pragma=busy_timeout(5000)&_pragma=journal_mode(WAL)&_pragma=synchronous(NORMAL)&_pragma=foreign_keys(ON)")
|
|
if err != nil {
|
|
return nil, fmt.Errorf("open db %q: %w", path, err)
|
|
}
|
|
if err := ensureSchema(db); err != nil {
|
|
db.Close()
|
|
return nil, err
|
|
}
|
|
return &Store{db: db}, nil
|
|
}
|
|
|
|
// Close releases the database handle.
|
|
func (s *Store) Close() error {
|
|
return s.db.Close()
|
|
}
|
|
|
|
func ensureSchema(db *sql.DB) error {
|
|
stmts := []string{
|
|
`CREATE TABLE IF NOT EXISTS requests (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
ts INTEGER NOT NULL,
|
|
model TEXT NOT NULL,
|
|
stream INTEGER NOT NULL DEFAULT 0,
|
|
prompt_tokens INTEGER NOT NULL DEFAULT 0,
|
|
completion_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
reasoning_tokens INTEGER NOT NULL DEFAULT 0,
|
|
cached_tokens INTEGER NOT NULL DEFAULT 0,
|
|
status TEXT NOT NULL DEFAULT 'success',
|
|
latency_ms INTEGER NOT NULL DEFAULT 0
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS credentials (
|
|
id INTEGER PRIMARY KEY CHECK (id = 1),
|
|
access_key TEXT NOT NULL DEFAULT '',
|
|
secret_key TEXT NOT NULL DEFAULT '',
|
|
token TEXT NOT NULL DEFAULT '',
|
|
api_key TEXT NOT NULL DEFAULT '',
|
|
model_base_url TEXT NOT NULL DEFAULT '',
|
|
email TEXT NOT NULL DEFAULT '',
|
|
organization TEXT NOT NULL DEFAULT '',
|
|
team TEXT NOT NULL DEFAULT '',
|
|
base_url TEXT NOT NULL DEFAULT '',
|
|
saved_at INTEGER NOT NULL DEFAULT 0
|
|
)`,
|
|
`CREATE INDEX IF NOT EXISTS idx_requests_ts ON requests(ts)`,
|
|
`CREATE INDEX IF NOT EXISTS idx_requests_model ON requests(model)`,
|
|
// Ensure the single credentials row exists so UPSERTs and SELECTs always
|
|
// have a target.
|
|
`INSERT INTO credentials (id) VALUES (1) ON CONFLICT(id) DO NOTHING`,
|
|
}
|
|
for _, q := range stmts {
|
|
if _, err := db.Exec(q); err != nil {
|
|
return fmt.Errorf("schema: %w", err)
|
|
}
|
|
}
|
|
// Add cached_tokens to databases created before this column existed. SQLite
|
|
// returns "duplicate column name" when it already exists; that is expected
|
|
// and ignored.
|
|
if _, err := db.Exec(`ALTER TABLE requests ADD COLUMN cached_tokens INTEGER NOT NULL DEFAULT 0`); err != nil {
|
|
if !strings.Contains(err.Error(), "duplicate column") {
|
|
return fmt.Errorf("migrate cached_tokens: %w", err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|