refactor: improve command flag registration

cmd/check.go

@@ -52,8 +52,6 @@ type channelExport struct {
 	Files     []exportFile `json:"files"`
 }
 
-const dateLayout = "2006-01-02_15-04-05"
-
 var termWidth = func() (width int, err error) {
 	width, _, err = term.GetSize(int(os.Stdout.Fd()))
 	if err == nil {
@@ -65,18 +63,15 @@ var termWidth = func() (width int, err error) {
 
 func NewCheckCmd() *cobra.Command {
 	var cfg config.ServerCmdConfig
+	loader := config.NewConfigLoader()
 	cmd := &cobra.Command{
 		Use:   "check",
 		Short: "Check and purge incomplete files",
 		Run: func(cmd *cobra.Command, args []string) {
 			runCheckCmd(cmd, &cfg)
 		},
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
-			loader := config.NewConfigLoader()
-			if err := loader.InitializeConfig(cmd); err != nil {
-				return err
-			}
-			if err := loader.Load(&cfg); err != nil {
+			if err := loader.Load(cmd, &cfg); err != nil {
 				return err
 			}
 			if err := checkRequiredCheckFlags(&cfg); err != nil {
@@ -85,18 +80,13 @@ func NewCheckCmd() *cobra.Command {
 			return nil
 		},
 	}
-	addChecktFlags(cmd, &cfg)
+	loader.RegisterPlags(cmd.Flags(), "", cfg, true)
+	cmd.Flags().Bool("export", true, "Export incomplete files to json file")
+	cmd.Flags().Bool("clean", false, "Clean missing and orphan file parts")
+	cmd.Flags().String("user", "", "Telegram User Name")
 	return cmd
 }
 
-func addChecktFlags(cmd *cobra.Command, cfg *config.ServerCmdConfig) {
-	flags := cmd.Flags()
-	config.AddCommonFlags(flags, cfg)
-	flags.Bool("export", true, "Export incomplete files to json file")
-	flags.Bool("clean", false, "Clean missing and orphan file parts")
-	flags.String("user", "", "Telegram User Name")
-}
-
 func checkRequiredCheckFlags(cfg *config.ServerCmdConfig) error {
 	var missingFields []string
 
@@ -245,6 +235,9 @@ func runCheckCmd(cmd *cobra.Command, cfg *config.ServerCmdConfig) {
 			lg.Fatalf("Channel %d: found %d messages out of %d", id, len(msgs), total)
 			continue
 		}
+
+		msgIds := utils.Map(msgs, func(m tg.NotEmptyMessage) int { return m.GetID() })
+
 		uploadPartIds := []int{}
 		if err := db.Model(&models.Upload{}).Where("user_id = ?", user.UserId).Where("channel_id = ?", id).
 			Pluck("part_id", &uploadPartIds).Error; err != nil {
@@ -256,8 +249,9 @@ func runCheckCmd(cmd *cobra.Command, cfg *config.ServerCmdConfig) {
 		for _, partID := range uploadPartIds {
 			uploadPartMap[partID] = true
 		}
+
 		msgMap := make(map[int]bool)
-		for _, m := range msgs {
+		for _, m := range msgIds {
 			if m > 0 && !uploadPartMap[m] {
 				msgMap[m] = true
 			}
@@ -343,7 +337,7 @@ func runCheckCmd(cmd *cobra.Command, cfg *config.ServerCmdConfig) {
 
 }
 
-func loadChannelMessages(ctx context.Context, client *telegram.Client, channelId int64) (msgs []int, total int, err error) {
+func loadChannelMessages(ctx context.Context, client *telegram.Client, channelId int64) (msgs []tg.NotEmptyMessage, total int, err error) {
 	errChan := make(chan error, 1)
 
 	go func() {
@@ -392,7 +386,7 @@ func loadChannelMessages(ctx context.Context, client *telegram.Client, channelId
 
 			for msgiter.Next(ctx) {
 				msg := msgiter.Value()
-				msgs = append(msgs, msg.Msg.GetID())
+				msgs = append(msgs, msg.Msg)
 				count++
 				bar.Set(count)
 			}

cmd/run.go

@@ -6,7 +6,6 @@ import (
 	"net"
 	"net/http"
 	"regexp"
-	"strings"
 	"time"
 
 	"github.com/go-chi/chi/v5"
@@ -21,7 +20,6 @@ import (
 	"github.com/tgdrive/teldrive/internal/chizap"
 	"github.com/tgdrive/teldrive/internal/config"
 	"github.com/tgdrive/teldrive/internal/database"
-	"github.com/tgdrive/teldrive/internal/duration"
 	"github.com/tgdrive/teldrive/internal/logging"
 	"github.com/tgdrive/teldrive/internal/middleware"
 	"github.com/tgdrive/teldrive/internal/tgc"
@@ -37,6 +35,7 @@ import (
 
 func NewRun() *cobra.Command {
 	var cfg config.ServerCmdConfig
+	loader := config.NewConfigLoader()
 	cmd := &cobra.Command{
 		Use:   "run",
 		Short: "Start Teldrive Server",
@@ -45,96 +44,18 @@ func NewRun() *cobra.Command {
 
 		},
 		PersistentPreRunE: func(cmd *cobra.Command, args []string) error {
-			loader := config.NewConfigLoader()
-			if err := loader.InitializeConfig(cmd); err != nil {
+			if err := loader.Load(cmd, &cfg); err != nil {
 				return err
 			}
-			if err := loader.Load(&cfg); err != nil {
-				return err
-			}
-			if err := checkRequiredRunFlags(&cfg); err != nil {
+			if err := loader.Validate(); err != nil {
 				return err
 			}
 			return nil
 		},
 	}
-	addServerFlags(cmd, &cfg)
+	loader.RegisterPlags(cmd.Flags(), "", cfg, false)
 	return cmd
 }
-func addServerFlags(cmd *cobra.Command, cfg *config.ServerCmdConfig) {
-
-	flags := cmd.Flags()
-
-	config.AddCommonFlags(flags, cfg)
-
-	// Server config
-	flags.IntVarP(&cfg.Server.Port, "server-port", "p", 8080, "Server port")
-	duration.DurationVar(flags, &cfg.Server.GracefulShutdown, "server-graceful-shutdown", 10*time.Second, "Server graceful shutdown timeout")
-	flags.BoolVar(&cfg.Server.EnablePprof, "server-enable-pprof", false, "Enable Pprof Profiling")
-	duration.DurationVar(flags, &cfg.Server.ReadTimeout, "server-read-timeout", 1*time.Hour, "Server read timeout")
-	duration.DurationVar(flags, &cfg.Server.WriteTimeout, "server-write-timeout", 1*time.Hour, "Server write timeout")
-
-	// CronJobs config
-	flags.BoolVar(&cfg.CronJobs.Enable, "cronjobs-enable", true, "Run cron jobs")
-	duration.DurationVar(flags, &cfg.CronJobs.CleanFilesInterval, "cronjobs-clean-files-interval", 1*time.Hour, "Clean files interval")
-	duration.DurationVar(flags, &cfg.CronJobs.CleanUploadsInterval, "cronjobs-clean-uploads-interval", 12*time.Hour, "Clean uploads interval")
-	duration.DurationVar(flags, &cfg.CronJobs.FolderSizeInterval, "cronjobs-folder-size-interval", 2*time.Hour, "Folder size update  interval")
-
-	// Cache config
-	flags.IntVar(&cfg.Cache.MaxSize, "cache-max-size", 10*1024*1024, "Max Cache max size (memory)")
-	flags.StringVar(&cfg.Cache.RedisAddr, "cache-redis-addr", "", "Redis address")
-	flags.StringVar(&cfg.Cache.RedisPass, "cache-redis-pass", "", "Redis password")
-
-	// JWT config
-	flags.StringVar(&cfg.JWT.Secret, "jwt-secret", "", "JWT secret key")
-	duration.DurationVar(flags, &cfg.JWT.SessionTime, "jwt-session-time", (30*24)*time.Hour, "JWT session duration")
-	flags.StringSliceVar(&cfg.JWT.AllowedUsers, "jwt-allowed-users", []string{}, "Allowed users")
-
-	// Telegram config
-	flags.StringVar(&cfg.TG.StorageFile, "tg-storage-file", "", "Sqlite Storage file path")
-	flags.BoolVar(&cfg.TG.RateLimit, "tg-rate-limit", true, "Enable rate limiting for telegram client")
-	flags.IntVar(&cfg.TG.RateBurst, "tg-rate-burst", 5, "Limiting burst for telegram client")
-	flags.IntVar(&cfg.TG.Rate, "tg-rate", 100, "Limiting rate for telegram client")
-	flags.StringVar(&cfg.TG.Proxy, "tg-proxy", "", "HTTP OR SOCKS5 proxy URL")
-	flags.BoolVar(&cfg.TG.DisableStreamBots, "tg-disable-stream-bots", false, "Disable Stream bots")
-	flags.BoolVar(&cfg.TG.Ntp, "tg-ntp", false, "Use NTP server time")
-	flags.BoolVar(&cfg.TG.EnableLogging, "tg-enable-logging", false, "Enable telegram client logging")
-	flags.Int64Var(&cfg.TG.PoolSize, "tg-pool-size", 8, "Telegram Session pool size")
-
-	// Telegram Uploads config
-	flags.StringVar(&cfg.TG.Uploads.EncryptionKey, "tg-uploads-encryption-key", "", "Uploads encryption key")
-	flags.IntVar(&cfg.TG.Uploads.Threads, "tg-uploads-threads", 8, "Uploads threads")
-	flags.IntVar(&cfg.TG.Uploads.MaxRetries, "tg-uploads-max-retries", 10, "Uploads Retries")
-	duration.DurationVar(flags, &cfg.TG.ReconnectTimeout, "tg-reconnect-timeout", 5*time.Minute, "Reconnect Timeout")
-	duration.DurationVar(flags, &cfg.TG.Uploads.Retention, "tg-uploads-retention", (24*7)*time.Hour, "Uploads retention duration")
-	flags.IntVar(&cfg.TG.Stream.MultiThreads, "tg-stream-multi-threads", 0, "Stream multi-threads")
-	flags.IntVar(&cfg.TG.Stream.Buffers, "tg-stream-buffers", 8, "No of Stream buffers")
-	duration.DurationVar(flags, &cfg.TG.Stream.ChunkTimeout, "tg-stream-chunk-timeout", 20*time.Second, "Chunk Fetch Timeout")
-
-}
-
-func checkRequiredRunFlags(cfg *config.ServerCmdConfig) error {
-	var missingFields []string
-
-	if cfg.DB.DataSource == "" {
-		missingFields = append(missingFields, "db-data-source")
-	}
-	if cfg.JWT.Secret == "" {
-		missingFields = append(missingFields, "jwt-secret")
-	}
-	if cfg.TG.AppHash == "" {
-		missingFields = append(missingFields, "tg-app-hash")
-	}
-	if cfg.TG.AppId == 0 {
-		missingFields = append(missingFields, "tg-app-id")
-	}
-
-	if len(missingFields) > 0 {
-		return fmt.Errorf("required configuration values not set: %s", strings.Join(missingFields, ", "))
-	}
-
-	return nil
-}
 
 func findAvailablePort(startPort int) (int, error) {
 	for port := startPort; port < startPort+100; port++ {

config.sample.toml

@@ -1,53 +1,57 @@
-[db]
-  data-source = ""
-  prepare-stmt = true
-  log-level = 1
-  [db.pool]
-    enable = true
-    max-idle-connections = 25
-    max-lifetime = "10m"
-    max-open-connections = 25
+[cache]
+max-size = 10485760
+redis-addr = ''
+redis-pass = ''
 
 [cronjobs]
-  enable = true
+clean-files-interval = '1h'
+clean-uploads-interval = '12h'
+enable = true
+folder-size-interval = '2h'
+
+[db]
+log-level = 'info'
+prepare-stmt = true
+
+[db.pool]
+enable = true
+max-idle-connections = 25
+max-lifetime = '10m'
+max-open-connections = 25
 
 [jwt]
-  allowed-users = [""]
-  secret = ""
-  session-time = "30d"
+session-time = '30d'
+secret = ''
+allowed-users = []
 
 [log]
-  development = true
-  level = -1
+level = 'info'
+file = ''
 
 [server]
-  graceful-shutdown = "15s"
-  port = 8080
-  read-timeout = "1h"
-  write-timeout = "1h"
+graceful-shutdown = '10s'
+port = 8080
+read-timeout = '1h'
+write-timeout = '1h'
 
 [tg]
-  app-hash = ""
-  app-id = 0
-  app-version = "4.6.3 K"
-  bg-bots-limit = 5
-  device-model = "Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:109.0) Gecko/20100101 Firefox/116.0"
-  disable-stream-bots = false
-  lang-code = "en"
-  lang-pack = "webk"
-  rate = 100
-  rate-burst = 5
-  rate-limit = true
-  session-file = ""
-  system-lang-code = "en-US"
-  system-version = "Win32"
-  proxy= "http://127.0.0.1:8080"
-
-  [tg.uploads]
-    encryption-key = ""
-    retention = "7d"
-    threads = 8
-  [tg.stream]
-    multi-threads = 0
-    buffers = 8
+pool-size = 8
+rate = 100
+rate-burst = 5
+ntp = false
+disable-stream-bots = false
+storage-file = ''
+proxy = ''
+rate-limit = true
+reconnect-timeout = '5m'
+
+[tg.stream]
+buffers = 8
+chunk-timeout = '20s'
 
+[tg.uploads]
+multi-threads = 0
+encryption-key = ''
+max-retries = 10
+retention = '7d'
+threads = '8'

go.mod

@@ -10,14 +10,14 @@ require (
 	github.com/go-chi/chi/v5 v5.2.0
 	github.com/go-chi/cors v1.2.1
 	github.com/go-co-op/gocron v1.37.0
-	github.com/go-viper/mapstructure/v2 v2.2.1
 	github.com/golang-jwt/jwt/v5 v5.2.1
 	github.com/google/uuid v1.6.0
 	github.com/gotd/contrib v0.21.0
 	github.com/gotd/td v0.117.0
 	github.com/iyear/connectproxy v0.1.1
 	github.com/k0kubun/go-ansi v0.0.0-20180517002512-3bf9e2903213
 	github.com/manifoldco/promptui v0.9.0
+	github.com/mitchellh/mapstructure v1.5.0
 	github.com/ogen-go/ogen v1.8.1
 	github.com/redis/go-redis/v9 v9.7.0
 	github.com/schollz/progressbar/v3 v3.18.0
@@ -58,7 +58,6 @@ require (
 	github.com/mattn/go-sqlite3 v1.14.24 // indirect
 	github.com/mfridman/interpolate v0.0.2 // indirect
 	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
-	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/ncruces/go-strftime v0.1.9 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect

go.sum

@@ -93,8 +93,6 @@ github.com/go-logr/stdr v1.2.2/go.mod h1:mMo/vtBO5dYbehREoey6XUKy/eSumjCCveDpRre
 github.com/go-sql-driver/mysql v1.7.0/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
 github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
 github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
-github.com/go-viper/mapstructure/v2 v2.2.1 h1:ZAaOCxANMuZx5RCeg0mBdEZk7DZasvvZIxtHqx8aGss=
-github.com/go-viper/mapstructure/v2 v2.2.1/go.mod h1:oJDH3BJKyqBA2TXFhDsKDGDTlndYOZ6rGS0BRZIxGhM=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
 github.com/golang-jwt/jwt/v5 v5.2.1 h1:OuVbFODueb089Lh128TAcimifWaLhJwVflnrgM17wHk=