ez-api/internal/config/config.go

package config

import (
	"fmt"
	"os"
	"strings"
	"time"

	"github.com/spf13/viper"
)

type Config struct {
	Server        ServerConfig
	Postgres      PostgresConfig
	Redis         RedisConfig
	Log           LogConfig
	Auth          AuthConfig
	ModelRegistry ModelRegistryConfig
}

type ServerConfig struct {
	Port string
}

type AuthConfig struct {
	JWTSecret string
}

type PostgresConfig struct {
	DSN string
}

type RedisConfig struct {
	Addr     string
	Password string
	DB       int
}

type LogConfig struct {
	BatchSize     int
	FlushInterval time.Duration
	QueueCapacity int
}

type ModelRegistryConfig struct {
	Enabled          bool
	RefreshSeconds   int
	ModelsDevBaseURL string
	ModelsDevRef     string
	CacheDir         string
	TimeoutSeconds   int
}

func Load() (*Config, error) {
	v := viper.New()

	v.SetDefault("server.port", "8080")
	v.SetDefault("postgres.dsn", "host=localhost user=postgres password=postgres dbname=ezapi port=5432 sslmode=disable")
	v.SetDefault("redis.addr", "localhost:6379")
	v.SetDefault("redis.password", "")
	v.SetDefault("redis.db", 0)
	v.SetDefault("log.batch_size", 10)
	v.SetDefault("log.flush_ms", 1000)
	v.SetDefault("log.queue_capacity", 10000)
	v.SetDefault("auth.jwt_secret", "change_me_in_production")
	v.SetDefault("model_registry.enabled", false)
	v.SetDefault("model_registry.refresh_seconds", 1800)
	v.SetDefault("model_registry.models_dev_base_url", "https://codeload.github.com/sst/models.dev/tar.gz")
	v.SetDefault("model_registry.models_dev_ref", "dev")
	v.SetDefault("model_registry.cache_dir", "./data/model-registry")
	v.SetDefault("model_registry.timeout_seconds", 30)

	v.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
	v.AutomaticEnv()

	_ = v.BindEnv("server.port", "EZ_API_PORT")
	_ = v.BindEnv("postgres.dsn", "EZ_PG_DSN")
	_ = v.BindEnv("redis.addr", "EZ_REDIS_ADDR")
	_ = v.BindEnv("redis.password", "EZ_REDIS_PASSWORD")
	_ = v.BindEnv("redis.db", "EZ_REDIS_DB")
	_ = v.BindEnv("log.batch_size", "EZ_LOG_BATCH_SIZE")
	_ = v.BindEnv("log.flush_ms", "EZ_LOG_FLUSH_MS")
	_ = v.BindEnv("log.queue_capacity", "EZ_LOG_QUEUE")
	_ = v.BindEnv("auth.jwt_secret", "EZ_JWT_SECRET")
	_ = v.BindEnv("model_registry.enabled", "EZ_MODEL_REGISTRY_ENABLED")
	_ = v.BindEnv("model_registry.refresh_seconds", "EZ_MODEL_REGISTRY_REFRESH_SECONDS")
	_ = v.BindEnv("model_registry.models_dev_base_url", "EZ_MODEL_REGISTRY_MODELS_DEV_BASE_URL")
	_ = v.BindEnv("model_registry.models_dev_ref", "EZ_MODEL_REGISTRY_MODELS_DEV_REF")
	_ = v.BindEnv("model_registry.cache_dir", "EZ_MODEL_REGISTRY_CACHE_DIR")
	_ = v.BindEnv("model_registry.timeout_seconds", "EZ_MODEL_REGISTRY_TIMEOUT_SECONDS")

	if configFile := os.Getenv("EZ_CONFIG_FILE"); configFile != "" {
		v.SetConfigFile(configFile)
	} else {
		v.SetConfigName("config")
		v.SetConfigType("yaml")
		v.AddConfigPath(".")
		v.AddConfigPath("./config")
	}

	if err := v.ReadInConfig(); err != nil {
		if _, ok := err.(viper.ConfigFileNotFoundError); !ok {
			return nil, fmt.Errorf("read config: %w", err)
		}
	}

	cfg := &Config{
		Server: ServerConfig{
			Port: v.GetString("server.port"),
		},
		Postgres: PostgresConfig{
			DSN: v.GetString("postgres.dsn"),
		},
		Redis: RedisConfig{
			Addr:     v.GetString("redis.addr"),
			Password: v.GetString("redis.password"),
			DB:       v.GetInt("redis.db"),
		},
		Log: LogConfig{
			BatchSize:     v.GetInt("log.batch_size"),
			FlushInterval: time.Duration(v.GetInt("log.flush_ms")) * time.Millisecond,
			QueueCapacity: v.GetInt("log.queue_capacity"),
		},
		Auth: AuthConfig{
			JWTSecret: v.GetString("auth.jwt_secret"),
		},
		ModelRegistry: ModelRegistryConfig{
			Enabled:          v.GetBool("model_registry.enabled"),
			RefreshSeconds:   v.GetInt("model_registry.refresh_seconds"),
			ModelsDevBaseURL: v.GetString("model_registry.models_dev_base_url"),
			ModelsDevRef:     v.GetString("model_registry.models_dev_ref"),
			CacheDir:         v.GetString("model_registry.cache_dir"),
			TimeoutSeconds:   v.GetInt("model_registry.timeout_seconds"),
		},
	}

	return cfg, nil
}