reforger_crawler_main/initializers/initialize.go

package initializers

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"time"

	"gitea.tbdevent.eu/TBD/reforger_crawler_main/models"
	"gopkg.in/yaml.v3"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
)

var DB *gorm.DB
var ServerPort string
var ServerHost string
var SecretKey string
var DiscordWebhookURL string
var AdminSecret string
var ScraperWebhookURL string
var DatabaseDriver string
var SQLitePath string
var PostgresDSN string

// Admin dashboard, metrics and telemetry settings
var TrustedProxies []string
var AdminSessionTTL time.Duration
var AdminCookieSecure string // auto|true|false
var MetricsEnabled bool
var MetricsToken string
var MetricsMinuteRetentionDays int
var MetricsHourRetentionDays int
var EventsRetentionDays int
var TopItemsRetentionDays int
var IndexTimeout time.Duration

func ConnectToDB() {
	var (
		db  *gorm.DB
		err error
	)

	switch strings.ToLower(DatabaseDriver) {
	case "", "sqlite", "sqlite3":
		sqlitePath := SQLitePath
		if sqlitePath == "" {
			sqlitePath = "crawler.db"
		}

		dsn := fmt.Sprintf("%s?_busy_timeout=5000&_journal_mode=WAL", sqlitePath)
		db, err = gorm.Open(sqlite.Open(dsn), &gorm.Config{
			SkipDefaultTransaction: true,
		})
		if err != nil {
			log.Fatalf("Failed to connect to sqlite database: %v", err)
		}

		if err := db.Exec("PRAGMA synchronous=NORMAL").Error; err != nil {
			log.Fatalf("Failed to set PRAGMA synchronous: %v", err)
		}
		if err := db.Exec("PRAGMA cache_size=-64000").Error; err != nil {
			log.Fatalf("Failed to set PRAGMA cache_size: %v", err)
		}
		if err := db.Exec("PRAGMA temp_store=MEMORY").Error; err != nil {
			log.Fatalf("Failed to set PRAGMA temp_store: %v", err)
		}
		if err := db.Exec("PRAGMA mmap_size=268435456").Error; err != nil {
			log.Fatalf("Failed to set PRAGMA mmap_size: %v", err)
		}
	case "postgres", "postgresql":
		if PostgresDSN == "" {
			log.Fatal("PostgreSQL selected but no DSN is configured")
		}

		db, err = gorm.Open(postgres.Open(PostgresDSN), &gorm.Config{
			SkipDefaultTransaction: true,
		})
		if err != nil {
			log.Fatalf("Failed to connect to postgres database: %v", err)
		}
	default:
		log.Fatalf("Unsupported database driver: %s", DatabaseDriver)
	}

	DB = db

	if err := DB.AutoMigrate(
		&models.Addon{}, &models.AddonFile{}, &models.WhitelistedHash{}, &models.AddonDependency{},
		&models.MetricSeries{}, &models.MetricPoint{}, &models.ServiceInstance{}, &models.ServiceEvent{}, &models.DailyTopItem{},
	); err != nil {
		log.Fatalf("Failed to run migrations: %v", err)
	}

	if err := DB.Exec(`CREATE INDEX IF NOT EXISTS idx_addon_files_hash_addon_optimized 
		ON addon_files(hash, addon_id) WHERE deleted_at IS NULL`).Error; err != nil {
		log.Fatalf("Failed to create idx_addon_files_hash_addon_optimized: %v", err)
	}

	if err := DB.Exec(`CREATE INDEX IF NOT EXISTS idx_whitelisted_hash 
		ON whitelisted_hashes(hash)`).Error; err != nil {
		log.Fatalf("Failed to create idx_whitelisted_hash: %v", err)
	}

	if err := DB.Exec(`CREATE INDEX IF NOT EXISTS idx_addons_id 
		ON addons(id)`).Error; err != nil {
		log.Fatalf("Failed to create idx_addons_id: %v", err)
	}

	if err := DB.Exec(`CREATE INDEX IF NOT EXISTS idx_addon_files_covering 
		ON addon_files(addon_id, hash, path, version) WHERE deleted_at IS NULL`).Error; err != nil {
		log.Fatalf("Failed to create idx_addon_files_covering: %v", err)
	}

	// Indexes used by the indexing queue and the admin dashboard; failures are
	// not fatal - they only cost performance
	optionalIndexes := []string{
		`CREATE INDEX IF NOT EXISTS idx_addons_queue ON addons(priority_indexing, updated_at)
			WHERE to_be_indexed = true AND is_being_indexed = false AND deleted_at IS NULL`,
		`CREATE INDEX IF NOT EXISTS idx_addons_created_at ON addons(created_at)`,
		`CREATE INDEX IF NOT EXISTS idx_addons_subscriber_count ON addons(subscriber_count)`,
		`CREATE INDEX IF NOT EXISTS idx_addons_current_version_size ON addons(current_version_size)`,
	}
	for _, stmt := range optionalIndexes {
		if err := DB.Exec(stmt).Error; err != nil {
			log.Printf("Warning: failed to create index: %v", err)
		}
	}
}

type Configuration struct {
	ServerPort        string `yaml:"server_port"`
	ServerHost        string `yaml:"server_host"`
	SecretKey         string `yaml:"secret_key"`
	AdminSecret       string `yaml:"admin_secret"`
	DiscordWebhookURL string `yaml:"discord_webhook_url"`
	ScraperWebhookURL string `yaml:"scraper_webhook_url"`
	DatabaseDriver    string `yaml:"db_driver"`
	SQLitePath        string `yaml:"sqlite_path"`
	PostgresDSN       string `yaml:"postgres_dsn"`
	PostgresHost      string `yaml:"postgres_host"`
	PostgresPort      string `yaml:"postgres_port"`
	PostgresUser      string `yaml:"postgres_user"`
	PostgresPassword  string `yaml:"postgres_password"`
	PostgresDBName    string `yaml:"postgres_db_name"`
	PostgresSSLMode   string `yaml:"postgres_sslmode"`
	TrustedProxies    string `yaml:"trusted_proxies"`
	AdminSessionTTLH  int    `yaml:"admin_session_ttl_hours"`
	AdminCookieSecure string `yaml:"admin_cookie_secure"`
	MetricsEnabled    *bool  `yaml:"metrics_enabled"`
	MetricsToken      string `yaml:"metrics_token"`
	MetricsMinuteDays int    `yaml:"metrics_minute_retention_days"`
	MetricsHourDays   int    `yaml:"metrics_hour_retention_days"`
	EventsDays        int    `yaml:"events_retention_days"`
	TopItemsDays      int    `yaml:"top_items_retention_days"`
	IndexTimeoutMin   int    `yaml:"index_timeout_minutes"`
	LegacyPort        string `yaml:"port"`
	LegacyIP          string `yaml:"ip"`
	LegacySecret      string `yaml:"secret"`
	LegacyDB          string `yaml:"db"`
	LegacyDiscordHook string `yaml:"discordWebhook"`
	LegacyAdminSecret string `yaml:"adminSecret"`
	LegacyScraperHook string `yaml:"scraperWebhook"`
}

func Load() {
	configuration := Configuration{
		ServerPort:      "8083",
		ServerHost:      "0.0.0.0",
		DatabaseDriver:  "sqlite",
		SQLitePath:      "crawler.db",
		PostgresSSLMode: "disable",

		AdminSessionTTLH:  168,
		AdminCookieSecure: "auto",
		MetricsMinuteDays: 7,
		MetricsHourDays:   180,
		EventsDays:        30,
		TopItemsDays:      90,
		IndexTimeoutMin:   30,
	}

	file, err := os.ReadFile("config.yaml")
	if err == nil {
		if err := yaml.Unmarshal(file, &configuration); err != nil {
			log.Fatalf("Failed to read yaml file: %v", err)
		}
	} else if !errors.Is(err, os.ErrNotExist) {
		log.Fatalf("Failed to open config file: %v", err)
	}

	applyLegacyConfigFallbacks(&configuration)
	applyEnvOverrides(&configuration)

	ServerPort = configuration.ServerPort
	ServerHost = configuration.ServerHost
	SecretKey = configuration.SecretKey
	AdminSecret = configuration.AdminSecret
	DiscordWebhookURL = configuration.DiscordWebhookURL
	ScraperWebhookURL = configuration.ScraperWebhookURL
	DatabaseDriver = strings.ToLower(configuration.DatabaseDriver)
	SQLitePath = configuration.SQLitePath

	TrustedProxies = splitList(configuration.TrustedProxies)
	AdminSessionTTL = time.Duration(positiveOr(configuration.AdminSessionTTLH, 168)) * time.Hour
	AdminCookieSecure = strings.ToLower(configuration.AdminCookieSecure)
	MetricsEnabled = configuration.MetricsEnabled == nil || *configuration.MetricsEnabled
	MetricsToken = configuration.MetricsToken
	MetricsMinuteRetentionDays = positiveOr(configuration.MetricsMinuteDays, 7)
	MetricsHourRetentionDays = positiveOr(configuration.MetricsHourDays, 180)
	EventsRetentionDays = positiveOr(configuration.EventsDays, 30)
	TopItemsRetentionDays = positiveOr(configuration.TopItemsDays, 90)
	IndexTimeout = time.Duration(positiveOr(configuration.IndexTimeoutMin, 30)) * time.Minute

	if strings.TrimSpace(configuration.PostgresDSN) != "" {
		PostgresDSN = configuration.PostgresDSN
	} else {
		PostgresDSN = buildPostgresDSN(configuration)
	}
}

func applyLegacyConfigFallbacks(c *Configuration) {
	if c.ServerPort == "" {
		c.ServerPort = c.LegacyPort
	}
	if c.ServerHost == "" {
		c.ServerHost = c.LegacyIP
	}
	if c.SecretKey == "" {
		c.SecretKey = c.LegacySecret
	}
	if c.AdminSecret == "" {
		c.AdminSecret = c.LegacyAdminSecret
	}
	if c.DiscordWebhookURL == "" {
		c.DiscordWebhookURL = c.LegacyDiscordHook
	}
	if c.ScraperWebhookURL == "" {
		c.ScraperWebhookURL = c.LegacyScraperHook
	}
	if c.SQLitePath == "" {
		c.SQLitePath = c.LegacyDB
	}
}

func applyEnvOverrides(c *Configuration) {
	c.ServerPort = firstEnv([]string{"SERVER_PORT", "PORT"}, c.ServerPort)
	c.ServerHost = firstEnv([]string{"SERVER_HOST", "IP"}, c.ServerHost)
	c.SecretKey = firstEnv([]string{"SECRET_KEY", "SECRET"}, c.SecretKey)
	c.AdminSecret = firstEnv([]string{"ADMIN_SECRET"}, c.AdminSecret)
	c.DiscordWebhookURL = firstEnv([]string{"DISCORD_WEBHOOK_URL", "DISCORD_WEBHOOK"}, c.DiscordWebhookURL)
	c.ScraperWebhookURL = firstEnv([]string{"SCRAPER_WEBHOOK_URL", "SCRAPER_WEBHOOK"}, c.ScraperWebhookURL)
	c.DatabaseDriver = firstEnv([]string{"DB_DRIVER"}, c.DatabaseDriver)
	c.SQLitePath = firstEnv([]string{"SQLITE_PATH", "DB"}, c.SQLitePath)
	c.PostgresDSN = firstEnv([]string{"POSTGRES_DSN"}, c.PostgresDSN)
	c.PostgresHost = firstEnv([]string{"POSTGRES_HOST"}, c.PostgresHost)
	c.PostgresPort = firstEnv([]string{"POSTGRES_PORT"}, c.PostgresPort)
	c.PostgresUser = firstEnv([]string{"POSTGRES_USER"}, c.PostgresUser)
	c.PostgresPassword = firstEnv([]string{"POSTGRES_PASSWORD"}, c.PostgresPassword)
	c.PostgresDBName = firstEnv([]string{"POSTGRES_DB_NAME"}, c.PostgresDBName)
	c.PostgresSSLMode = firstEnv([]string{"POSTGRES_SSLMODE"}, c.PostgresSSLMode)
	c.TrustedProxies = firstEnv([]string{"TRUSTED_PROXIES"}, c.TrustedProxies)
	c.AdminSessionTTLH = firstEnvInt([]string{"ADMIN_SESSION_TTL_HOURS"}, c.AdminSessionTTLH)
	c.AdminCookieSecure = firstEnv([]string{"ADMIN_COOKIE_SECURE"}, c.AdminCookieSecure)
	if v := firstEnv([]string{"METRICS_ENABLED"}, ""); v != "" {
		enabled := v == "1" || strings.EqualFold(v, "true") || strings.EqualFold(v, "yes")
		c.MetricsEnabled = &enabled
	}
	c.MetricsToken = firstEnv([]string{"METRICS_TOKEN"}, c.MetricsToken)
	c.MetricsMinuteDays = firstEnvInt([]string{"METRICS_MINUTE_RETENTION_DAYS"}, c.MetricsMinuteDays)
	c.MetricsHourDays = firstEnvInt([]string{"METRICS_HOUR_RETENTION_DAYS"}, c.MetricsHourDays)
	c.EventsDays = firstEnvInt([]string{"EVENTS_RETENTION_DAYS"}, c.EventsDays)
	c.TopItemsDays = firstEnvInt([]string{"TOP_ITEMS_RETENTION_DAYS"}, c.TopItemsDays)
	c.IndexTimeoutMin = firstEnvInt([]string{"INDEX_TIMEOUT_MINUTES"}, c.IndexTimeoutMin)
}

func firstEnvInt(keys []string, fallback int) int {
	raw := firstEnv(keys, "")
	if raw == "" {
		return fallback
	}
	val, err := strconv.Atoi(strings.TrimSpace(raw))
	if err != nil {
		log.Printf("Warning: ignoring invalid integer %q for %s", raw, strings.Join(keys, "/"))
		return fallback
	}
	return val
}

func positiveOr(val, fallback int) int {
	if val > 0 {
		return val
	}
	return fallback
}

func splitList(raw string) []string {
	var out []string
	for _, part := range strings.Split(raw, ",") {
		if part = strings.TrimSpace(part); part != "" {
			out = append(out, part)
		}
	}
	return out
}

func firstEnv(keys []string, fallback string) string {
	for _, k := range keys {
		if value, ok := os.LookupEnv(k); ok && strings.TrimSpace(value) != "" {
			return value
		}
	}
	return fallback
}

func buildPostgresDSN(c Configuration) string {
	if c.PostgresHost == "" || c.PostgresUser == "" || c.PostgresDBName == "" {
		return ""
	}

	port := c.PostgresPort
	if port == "" {
		port = "5432"
	}

	sslMode := c.PostgresSSLMode
	if sslMode == "" {
		sslMode = "disable"
	}

	return fmt.Sprintf(
		"host=%s user=%s password=%s dbname=%s port=%s sslmode=%s",
		c.PostgresHost,
		c.PostgresUser,
		c.PostgresPassword,
		c.PostgresDBName,
		port,
		sslMode,
	)
}