feat: migrate backend from rust to go
This commit replaces the Axum/Rust backend with a Gin/Go implementation. The original Rust code has been archived in the 'rust' branch.
This commit is contained in:
264
internal/db/db.go
Normal file
264
internal/db/db.go
Normal file
@@ -0,0 +1,264 @@
|
||||
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("admin"), 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 'admin' (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"`
|
||||
}
|
||||
Reference in New Issue
Block a user