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"` Username string `db:"username"` PasswordHash string `db:"password_hash"` DisplayName *string `db:"display_name"` Role string `db:"role"` MustChangePassword bool `db:"must_change_password"` CreatedAt time.Time `db:"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"` }