Files
Luminary/internal/store/sqlite/store.go
T
renjue b1394a2817
CI / docker (push) Successful in 3m49s
Split ingress keys into text/image protocols with Replicate model listing.
Text keys use OpenAI routes; image keys use predictions with ingress Model ID,
OpenAPI input schemas, owner/name or short name resolution, protocol-aware admin
whitelist, and in-app API docs aligned to UI conventions.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-06-29 22:36:42 +08:00

120 lines
3.0 KiB
Go

package sqlite
import (
"fmt"
"os"
"path/filepath"
"github.com/glebarez/sqlite"
"github.com/rose_cat707/luminary/internal/model"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
type Store struct {
db *gorm.DB
}
func New(path string) (*Store, error) {
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
return nil, fmt.Errorf("create data dir: %w", err)
}
db, err := gorm.Open(sqlite.Open(path+"?_pragma=journal_mode(WAL)&_pragma=busy_timeout(5000)"), &gorm.Config{
Logger: logger.Default.LogMode(logger.Warn),
})
if err != nil {
return nil, fmt.Errorf("open sqlite: %w", err)
}
sqlDB, err := db.DB()
if err != nil {
return nil, err
}
sqlDB.SetMaxOpenConns(1)
if err := ensureIPRulesSchema(db); err != nil {
return nil, fmt.Errorf("migrate ip_rules: %w", err)
}
if err := db.AutoMigrate(
&model.Provider{},
&model.ProviderKey{},
&model.ModelEntry{},
&model.IngressKey{},
&model.AdminUser{},
&model.AdminSession{},
&model.UsageRecord{},
&model.HealthCheckRecord{},
&model.RoutingRule{},
&model.SystemSettings{},
&model.Prediction{},
&model.StoredImage{},
); err != nil {
return nil, fmt.Errorf("migrate: %w", err)
}
if err := ensureIngressKeyProtocol(db); err != nil {
return nil, fmt.Errorf("migrate ingress protocol: %w", err)
}
return &Store{db: db}, nil
}
func ensureIngressKeyProtocol(db *gorm.DB) error {
if !tableExists(db, "ingress_keys") {
return nil
}
return db.Exec(`UPDATE ingress_keys SET protocol = ? WHERE protocol IS NULL OR protocol = ''`, model.IngressProtocolText).Error
}
// ensureIPRulesSchema manages ip_rules without GORM AutoMigrate (avoids bad table rebuilds).
func ensureIPRulesSchema(db *gorm.DB) error {
if !tableExists(db, "ip_rules") {
return db.Exec(`
CREATE TABLE ip_rules (
id INTEGER PRIMARY KEY AUTOINCREMENT,
cidr TEXT NOT NULL,
type TEXT NOT NULL,
scope TEXT NOT NULL DEFAULT 'proxy',
enabled NUMERIC DEFAULT 1,
created_at DATETIME
);
`).Error
}
if !tableHasColumn(db, "ip_rules", "c_id_r") {
return nil
}
cidrExpr := "c_id_r"
if tableHasColumn(db, "ip_rules", "cidr") {
cidrExpr = "COALESCE(NULLIF(cidr, ''), c_id_r)"
}
return db.Exec(`
CREATE TABLE ip_rules_mig (
id INTEGER PRIMARY KEY AUTOINCREMENT,
cidr TEXT NOT NULL,
type TEXT NOT NULL,
scope TEXT NOT NULL DEFAULT 'proxy',
enabled NUMERIC DEFAULT 1,
created_at DATETIME
);
INSERT INTO ip_rules_mig (id, cidr, type, scope, enabled, created_at)
SELECT id, ` + cidrExpr + `, type, scope, enabled, created_at
FROM ip_rules;
DROP TABLE ip_rules;
ALTER TABLE ip_rules_mig RENAME TO ip_rules;
`).Error
}
func tableHasColumn(db *gorm.DB, table, column string) bool {
var n int64
db.Raw(`SELECT COUNT(*) FROM pragma_table_info(?) WHERE name = ?`, table, column).Scan(&n)
return n > 0
}
func tableExists(db *gorm.DB, table string) bool {
var n int64
db.Raw(`SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?`, table).Scan(&n)
return n > 0
}
func (s *Store) DB() *gorm.DB {
return s.db
}