gpt-load/internal/storage/database.go

package storage

import (
	"context"
	"database/sql"
	"fmt"
	"net/url"
	"strings"

	"github.com/glebarez/sqlite"
	"github.com/sirupsen/logrus"
	gormmysql "gorm.io/driver/mysql"
	gormpostgres "gorm.io/driver/postgres"
	"gorm.io/gorm"

	"gpt-load/internal/platform/config"
)

const (
	databaseMaxOpenConnections = 10
	databaseMaxIdleConnections = 5
)

// openDatabase is the shared GORM/SQL lifecycle for every supported driver.
// Driver-specific behavior is limited to dialector construction and the
// SQLite runtime hook in db.go.
func openDatabase(driver config.DatabaseDriver, dialector gorm.Dialector) (*gorm.DB, error) {
	db, err := gorm.Open(dialector, &gorm.Config{
		Logger:         databaseLogger,
		TranslateError: true,
	})
	if err != nil {
		return nil, fmt.Errorf("open %s database: %w", databaseDisplayName(driver), err)
	}

	sqlDB, err := db.DB()
	if err != nil {
		return nil, fmt.Errorf("get %s connection pool: %w", databaseDisplayName(driver), err)
	}
	configureDatabasePool(sqlDB, driver)
	if err := sqlDB.PingContext(context.Background()); err != nil {
		_ = sqlDB.Close()
		return nil, fmt.Errorf("ping %s database: %w", databaseDisplayName(driver), err)
	}
	return db, nil
}

func configureDatabasePool(sqlDB *sql.DB, driver config.DatabaseDriver) {
	maxOpenConnections := databaseMaxOpenConnections
	maxIdleConnections := databaseMaxIdleConnections
	if driver == config.DatabaseDriverSQLite {
		// SQLite's single-writer runtime and shared :memory: compatibility both
		// require one physical connection.
		maxOpenConnections = 1
		maxIdleConnections = 1
	}
	sqlDB.SetMaxOpenConns(maxOpenConnections)
	sqlDB.SetMaxIdleConns(maxIdleConnections)
}

func newDatabaseDialector(database config.DatabaseConfig) (gorm.Dialector, error) {
	switch database.Driver {
	case config.DatabaseDriverSQLite:
		return sqlite.Open(database.DSN), nil
	case config.DatabaseDriverMySQL:
		dsn, err := mysqlDSNFromURL(database.DSN)
		if err != nil {
			return nil, err
		}
		return gormmysql.Open(dsn), nil
	case config.DatabaseDriverPostgreSQL:
		// Schema migrations rename/rebuild tables while the process is running.
		// pgx's implicit statement cache otherwise can retain a result shape from
		// the legacy table and fail the first query against the rebuilt table.
		return gormpostgres.New(gormpostgres.Config{
			DSN:                  database.DSN,
			PreferSimpleProtocol: true,
		}), nil
	default:
		return nil, fmt.Errorf("unsupported database driver")
	}
}

func mysqlDSNFromURL(rawDSN string) (string, error) {
	parsed, err := url.Parse(rawDSN)
	if err != nil || !strings.EqualFold(parsed.Scheme, "mysql") || parsed.Host == "" {
		return "", fmt.Errorf("DATABASE_DSN has an invalid MySQL URL")
	}
	databaseName := strings.TrimPrefix(parsed.EscapedPath(), "/")
	if databaseName == "" || strings.Contains(strings.TrimPrefix(parsed.Path, "/"), "/") {
		return "", fmt.Errorf("DATABASE_DSN MySQL URL must include one database name")
	}
	if parsed.Fragment != "" {
		return "", fmt.Errorf("DATABASE_DSN MySQL URL must not include a fragment")
	}
	query, err := url.ParseQuery(parsed.RawQuery)
	if err != nil {
		return "", fmt.Errorf("DATABASE_DSN has an invalid MySQL query")
	}
	// Keep the application-visible behavior stable across MySQL installations:
	// parseTime is required for time-valued driver fields, clientFoundRows is
	// required by existing RowsAffected contracts, and utf8mb4/binary lets
	// connection literals represent exact identifiers. Schema migration 0004
	// enforces the corresponding binary identity on model_prices.model_id.
	query.Set("parseTime", "true")
	query.Set("clientFoundRows", "true")
	if query.Get("charset") == "" {
		query.Set("charset", "utf8mb4")
	}
	if query.Get("collation") == "" {
		query.Set("collation", "utf8mb4_bin")
	}

	credentials := ""
	if parsed.User != nil {
		credentials = parsed.User.Username()
		if password, ok := parsed.User.Password(); ok {
			credentials += ":" + password
		}
		credentials += "@"
	}
	return credentials + "tcp(" + parsed.Host + ")/" + databaseName + "?" + query.Encode(), nil
}

func logExternalDatabaseSource(driver config.DatabaseDriver) {
	logrus.WithFields(logrus.Fields{
		"database_source": config.DatabaseSourceExternal,
		"database_driver": driver,
	}).Info("Database storage is managed by the operator")
}

func databaseDisplayName(driver config.DatabaseDriver) string {
	switch driver {
	case config.DatabaseDriverSQLite:
		return "SQLite"
	case config.DatabaseDriverMySQL:
		return "MySQL"
	case config.DatabaseDriverPostgreSQL:
		return "PostgreSQL"
	default:
		return string(driver)
	}
}