Merge pull request #22 from tbphp/feat-db-sqlite

feat: Multi-DB Driver

internal/db/database.go

@@ -5,9 +5,13 @@ import (
 	"gpt-load/internal/types"
 	"log"
 	"os"
+	"path/filepath"
+	"strings"
 	"time"
 
+	"github.com/glebarez/sqlite"
 	"gorm.io/driver/mysql"
+	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
 )
@@ -16,7 +20,8 @@ var DB *gorm.DB
 
 func NewDB(configManager types.ConfigManager) (*gorm.DB, error) {
 	dbConfig := configManager.GetDatabaseConfig()
-	if dbConfig.DSN == "" {
+	dsn := dbConfig.DSN
+	if dsn == "" {
 		return nil, fmt.Errorf("DATABASE_DSN is not configured")
 	}
 
@@ -33,9 +38,32 @@ func NewDB(configManager types.ConfigManager) (*gorm.DB, error) {
 		)
 	}
 
+	var dialector gorm.Dialector
+	if strings.HasPrefix(dsn, "postgres://") || strings.HasPrefix(dsn, "postgresql://") {
+		dialector = postgres.New(postgres.Config{
+			DSN:                  dsn,
+			PreferSimpleProtocol: true,
+		})
+	} else if strings.Contains(dsn, "@tcp") {
+		if !strings.Contains(dsn, "parseTime") {
+			if strings.Contains(dsn, "?") {
+				dsn += "&parseTime=true"
+			} else {
+				dsn += "?parseTime=true"
+			}
+		}
+		dialector = mysql.Open(dsn)
+	} else {
+		if err := os.MkdirAll(filepath.Dir(dsn), 0755); err != nil {
+			return nil, fmt.Errorf("failed to create database directory: %w", err)
+		}
+		dialector = sqlite.Open(dsn + "?_busy_timeout=5000")
+	}
+
 	var err error
-	DB, err = gorm.Open(mysql.Open(dbConfig.DSN), &gorm.Config{
-		Logger: newLogger,
+	DB, err = gorm.Open(dialector, &gorm.Config{
+		Logger:      newLogger,
+		PrepareStmt: true,
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to connect to database: %w", err)
@@ -45,10 +73,9 @@ func NewDB(configManager types.ConfigManager) (*gorm.DB, error) {
 	if err != nil {
 		return nil, fmt.Errorf("failed to get sql.DB: %w", err)
 	}
-
-	// Set connection pool parameters
-	sqlDB.SetMaxIdleConns(10)
-	sqlDB.SetMaxOpenConns(100)
+	// Set connection pool parameters for all drivers
+	sqlDB.SetMaxIdleConns(50)
+	sqlDB.SetMaxOpenConns(500)
 	sqlDB.SetConnMaxLifetime(time.Hour)
 
 	return DB, nil