hobokenchicken
0f0486d8d4
Added JSON tags to the User struct to match frontend expectations and excluded sensitive fields. Updated session management to include and persist DisplayName. Unified user field names (using display_name) across backend, sessions, and frontend UI.
265 lines
9.4 KiB
Go
265 lines
9.4 KiB
Go
package db
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"time"
|
|
|
|
"github.com/jmoiron/sqlx"
|
|
_ "modernc.org/sqlite"
|
|
"golang.org/x/crypto/bcrypt"
|
|
)
|
|
|
|
type DB struct {
|
|
*sqlx.DB
|
|
}
|
|
|
|
func Init(path string) (*DB, error) {
|
|
// Ensure directory exists
|
|
dir := filepath.Dir(path)
|
|
if dir != "." && dir != "/" {
|
|
if err := os.MkdirAll(dir, 0755); err != nil {
|
|
return nil, fmt.Errorf("failed to create database directory: %w", err)
|
|
}
|
|
}
|
|
|
|
// Connect to SQLite
|
|
dsn := fmt.Sprintf("file:%s?_pragma=foreign_keys(1)", path)
|
|
db, err := sqlx.Connect("sqlite", dsn)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to connect to database: %w", err)
|
|
}
|
|
|
|
instance := &DB{db}
|
|
|
|
// Run migrations
|
|
if err := instance.RunMigrations(); err != nil {
|
|
return nil, fmt.Errorf("failed to run migrations: %w", err)
|
|
}
|
|
|
|
return instance, nil
|
|
}
|
|
|
|
func (db *DB) RunMigrations() error {
|
|
// Tables creation
|
|
queries := []string{
|
|
`CREATE TABLE IF NOT EXISTS clients (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
client_id TEXT UNIQUE NOT NULL,
|
|
name TEXT,
|
|
description TEXT,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
rate_limit_per_minute INTEGER DEFAULT 60,
|
|
total_requests INTEGER DEFAULT 0,
|
|
total_tokens INTEGER DEFAULT 0,
|
|
total_cost REAL DEFAULT 0.0
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS llm_requests (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
client_id TEXT,
|
|
provider TEXT,
|
|
model TEXT,
|
|
prompt_tokens INTEGER,
|
|
completion_tokens INTEGER,
|
|
reasoning_tokens INTEGER DEFAULT 0,
|
|
total_tokens INTEGER,
|
|
cost REAL,
|
|
has_images BOOLEAN DEFAULT FALSE,
|
|
status TEXT DEFAULT 'success',
|
|
error_message TEXT,
|
|
duration_ms INTEGER,
|
|
request_body TEXT,
|
|
response_body TEXT,
|
|
cache_read_tokens INTEGER DEFAULT 0,
|
|
cache_write_tokens INTEGER DEFAULT 0,
|
|
FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE SET NULL
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS provider_configs (
|
|
id TEXT PRIMARY KEY,
|
|
display_name TEXT NOT NULL,
|
|
enabled BOOLEAN DEFAULT TRUE,
|
|
base_url TEXT,
|
|
api_key TEXT,
|
|
credit_balance REAL DEFAULT 0.0,
|
|
low_credit_threshold REAL DEFAULT 5.0,
|
|
billing_mode TEXT,
|
|
api_key_encrypted BOOLEAN DEFAULT FALSE,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS model_configs (
|
|
id TEXT PRIMARY KEY,
|
|
provider_id TEXT NOT NULL,
|
|
display_name TEXT,
|
|
enabled BOOLEAN DEFAULT TRUE,
|
|
prompt_cost_per_m REAL,
|
|
completion_cost_per_m REAL,
|
|
cache_read_cost_per_m REAL,
|
|
cache_write_cost_per_m REAL,
|
|
mapping TEXT,
|
|
updated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (provider_id) REFERENCES provider_configs(id) ON DELETE CASCADE
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS users (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
username TEXT UNIQUE NOT NULL,
|
|
password_hash TEXT NOT NULL,
|
|
display_name TEXT,
|
|
role TEXT DEFAULT 'admin',
|
|
must_change_password BOOLEAN DEFAULT FALSE,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
|
)`,
|
|
`CREATE TABLE IF NOT EXISTS client_tokens (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
client_id TEXT NOT NULL,
|
|
token TEXT NOT NULL UNIQUE,
|
|
name TEXT DEFAULT 'default',
|
|
is_active BOOLEAN DEFAULT TRUE,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
last_used_at DATETIME,
|
|
FOREIGN KEY (client_id) REFERENCES clients(client_id) ON DELETE CASCADE
|
|
)`,
|
|
}
|
|
|
|
for _, q := range queries {
|
|
if _, err := db.Exec(q); err != nil {
|
|
return fmt.Errorf("migration failed for query [%s]: %w", q, err)
|
|
}
|
|
}
|
|
|
|
// Add indices
|
|
indices := []string{
|
|
"CREATE INDEX IF NOT EXISTS idx_clients_client_id ON clients(client_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_clients_created_at ON clients(created_at)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_timestamp ON llm_requests(timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_client_id ON llm_requests(client_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_provider ON llm_requests(provider)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_status ON llm_requests(status)",
|
|
"CREATE UNIQUE INDEX IF NOT EXISTS idx_client_tokens_token ON client_tokens(token)",
|
|
"CREATE INDEX IF NOT EXISTS idx_client_tokens_client_id ON client_tokens(client_id)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_client_timestamp ON llm_requests(client_id, timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_llm_requests_provider_timestamp ON llm_requests(provider, timestamp)",
|
|
"CREATE INDEX IF NOT EXISTS idx_model_configs_provider_id ON model_configs(provider_id)",
|
|
}
|
|
|
|
for _, idx := range indices {
|
|
if _, err := db.Exec(idx); err != nil {
|
|
return fmt.Errorf("failed to create index [%s]: %w", idx, err)
|
|
}
|
|
}
|
|
|
|
// Default admin user
|
|
var count int
|
|
if err := db.Get(&count, "SELECT COUNT(*) FROM users"); err != nil {
|
|
return fmt.Errorf("failed to count users: %w", err)
|
|
}
|
|
|
|
if count == 0 {
|
|
hash, err := bcrypt.GenerateFromPassword([]byte("admin123"), 12)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to hash default password: %w", err)
|
|
}
|
|
_, err = db.Exec("INSERT INTO users (username, password_hash, role, must_change_password) VALUES ('admin', ?, 'admin', 1)", string(hash))
|
|
if err != nil {
|
|
return fmt.Errorf("failed to insert default admin: %w", err)
|
|
}
|
|
log.Println("Created default admin user with password 'admin123' (must change on first login)")
|
|
}
|
|
|
|
// Default client
|
|
_, err := db.Exec(`INSERT OR IGNORE INTO clients (client_id, name, description)
|
|
VALUES ('default', 'Default Client', 'Default client for anonymous requests')`)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to insert default client: %w", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// Data models for DB tables
|
|
|
|
type Client struct {
|
|
ID int `db:"id"`
|
|
ClientID string `db:"client_id"`
|
|
Name *string `db:"name"`
|
|
Description *string `db:"description"`
|
|
CreatedAt time.Time `db:"created_at"`
|
|
UpdatedAt time.Time `db:"updated_at"`
|
|
IsActive bool `db:"is_active"`
|
|
RateLimitPerMinute int `db:"rate_limit_per_minute"`
|
|
TotalRequests int `db:"total_requests"`
|
|
TotalTokens int `db:"total_tokens"`
|
|
TotalCost float64 `db:"total_cost"`
|
|
}
|
|
|
|
type LLMRequest struct {
|
|
ID int `db:"id"`
|
|
Timestamp time.Time `db:"timestamp"`
|
|
ClientID *string `db:"client_id"`
|
|
Provider *string `db:"provider"`
|
|
Model *string `db:"model"`
|
|
PromptTokens *int `db:"prompt_tokens"`
|
|
CompletionTokens *int `db:"completion_tokens"`
|
|
ReasoningTokens int `db:"reasoning_tokens"`
|
|
TotalTokens *int `db:"total_tokens"`
|
|
Cost *float64 `db:"cost"`
|
|
HasImages bool `db:"has_images"`
|
|
Status string `db:"status"`
|
|
ErrorMessage *string `db:"error_message"`
|
|
DurationMS *int `db:"duration_ms"`
|
|
RequestBody *string `db:"request_body"`
|
|
ResponseBody *string `db:"response_body"`
|
|
CacheReadTokens int `db:"cache_read_tokens"`
|
|
CacheWriteTokens int `db:"cache_write_tokens"`
|
|
}
|
|
|
|
type ProviderConfig struct {
|
|
ID string `db:"id"`
|
|
DisplayName string `db:"display_name"`
|
|
Enabled bool `db:"enabled"`
|
|
BaseURL *string `db:"base_url"`
|
|
APIKey *string `db:"api_key"`
|
|
CreditBalance float64 `db:"credit_balance"`
|
|
LowCreditThreshold float64 `db:"low_credit_threshold"`
|
|
BillingMode *string `db:"billing_mode"`
|
|
APIKeyEncrypted bool `db:"api_key_encrypted"`
|
|
UpdatedAt time.Time `db:"updated_at"`
|
|
}
|
|
|
|
type ModelConfig struct {
|
|
ID string `db:"id"`
|
|
ProviderID string `db:"provider_id"`
|
|
DisplayName *string `db:"display_name"`
|
|
Enabled bool `db:"enabled"`
|
|
PromptCostPerM *float64 `db:"prompt_cost_per_m"`
|
|
CompletionCostPerM *float64 `db:"completion_cost_per_m"`
|
|
CacheReadCostPerM *float64 `db:"cache_read_cost_per_m"`
|
|
CacheWriteCostPerM *float64 `db:"cache_write_cost_per_m"`
|
|
Mapping *string `db:"mapping"`
|
|
UpdatedAt time.Time `db:"updated_at"`
|
|
}
|
|
|
|
type User struct {
|
|
ID int `db:"id" json:"id"`
|
|
Username string `db:"username" json:"username"`
|
|
PasswordHash string `db:"password_hash" json:"-"`
|
|
DisplayName *string `db:"display_name" json:"display_name"`
|
|
Role string `db:"role" json:"role"`
|
|
MustChangePassword bool `db:"must_change_password" json:"must_change_password"`
|
|
CreatedAt time.Time `db:"created_at" json:"created_at"`
|
|
}
|
|
|
|
type ClientToken struct {
|
|
ID int `db:"id"`
|
|
ClientID string `db:"client_id"`
|
|
Token string `db:"token"`
|
|
Name string `db:"name"`
|
|
IsActive bool `db:"is_active"`
|
|
CreatedAt time.Time `db:"created_at"`
|
|
LastUsedAt *time.Time `db:"last_used_at"`
|
|
}
|