config: Allow overriding some config from command line

cmd/tidb-lightning-ctl/main.go

@@ -39,10 +39,17 @@ func main() {
 
 func run() error {
 	cfg := config.NewConfig()
-	cfg.FlagSet = flag.NewFlagSet("lightning-ctl", flag.ExitOnError)
-	fs := cfg.FlagSet
+	fs := flag.NewFlagSet("lightning-ctl", flag.ExitOnError)
 
-	fs.StringVar(&cfg.ConfigFile, "config", "tidb-lightning.toml", "tidb-lightning configuration file")
+	fs.StringVar(&cfg.ConfigFile, "config", "", "tidb-lightning configuration file")
+
+	logLevel := fs.String("L", "", `log level: info, debug, warn, error, fatal (default "info")`)
+	logFilePath := fs.String("log-file", "", "log file path")
+	tidbHost := fs.String("tidb-host", "", "TiDB server host")
+	tidbPort := fs.Int("tidb-port", 0, "TiDB server port (default 4000)")
+	tidbUser := fs.String("tidb-user", "", "TiDB user name to connect")
+	pdAddr := fs.String("pd-urls", "", "PD endpoint address")
+	importerAddr := fs.String("importer", "", "address (host:port) to connect to tikv-importer")
 
 	compact := fs.Bool("compact", false, "do manual compaction on the target cluster")
 	mode := fs.String("switch-mode", "", "switch tikv into import mode or normal mode, values can be ['import', 'normal']")
@@ -63,6 +70,30 @@ func run() error {
 		return errors.Trace(err)
 	}
 
+	if *logLevel != "" {
+		cfg.App.LogConfig.Level = *logLevel
+	}
+	if *logFilePath != "" {
+		cfg.App.LogConfig.File = *logFilePath
+	}
+	if *tidbHost != "" {
+		cfg.TiDB.Host = *tidbHost
+	}
+	if *tidbPort != 0 {
+		cfg.TiDB.Port = *tidbPort
+	}
+	if *tidbUser != "" {
+		cfg.TiDB.User = *tidbUser
+	}
+	if *pdAddr != "" {
+		cfg.TiDB.PdAddr = *pdAddr
+	}
+	if *importerAddr != "" {
+		cfg.TikvImporter.Addr = *importerAddr
+	}
+
+	cfg.Adjust()
+
 	ctx := context.Background()
 
 	if *compact {

lightning/config/config.go

@@ -35,6 +35,8 @@ const (
 	NormalMode = "normal"
 )
 
+var defaultConfigPaths = []string{"tidb-lightning.toml", "conf/tidb-lightning.toml"}
+
 type DBStore struct {
 	Host       string `toml:"host" json:"host"`
 	Port       int    `toml:"port" json:"port"`
@@ -54,8 +56,6 @@ type DBStore struct {
 }
 
 type Config struct {
-	*flag.FlagSet `json:"-"`
-
 	App  Lightning `toml:"lightning" json:"lightning"`
 	TiDB DBStore   `toml:"tidb" json:"tidb"`
 
@@ -69,10 +69,7 @@ type Config struct {
 	Cron         Cron            `toml:"cron" json:"cron"`
 
 	// command line flags
-	ConfigFile   string `json:"config-file"`
-	DoCompact    bool   `json:"-"`
-	SwitchMode   string `json:"-"`
-	printVersion bool
+	ConfigFile string `json:"config-file"`
 }
 
 func (c *Config) String() string {
@@ -164,6 +161,12 @@ func NewConfig() *Config {
 			CheckRequirements: true,
 		},
 		TiDB: DBStore{
+			Host:                       "127.0.0.1",
+			Port:                       4000,
+			User:                       "root",
+			StatusPort:                 10080,
+			PdAddr:                     "127.0.0.1:2379",
+			LogLevel:                   "error",
 			StrSQLMode:                 mysql.DefaultSQLMode,
 			BuildStatsConcurrency:      20,
 			DistSQLScanConcurrency:     100,
@@ -179,38 +182,83 @@ func NewConfig() *Config {
 				Separator: ",",
 			},
 		},
+		PostRestore: PostRestore{
+			Checksum: true,
+		},
 	}
 }
 
 func LoadConfig(args []string) (*Config, error) {
 	cfg := NewConfig()
 
-	cfg.FlagSet = flag.NewFlagSet("lightning", flag.ContinueOnError)
-	fs := cfg.FlagSet
+	fs := flag.NewFlagSet("lightning", flag.ContinueOnError)
 
 	// if both `-c` and `-config` are specified, the last one in the command line will take effect.
 	// the default value is assigned immediately after the StringVar() call,
 	// so it is fine to not give any default value for `-c`, to keep the `-h` page clean.
 	fs.StringVar(&cfg.ConfigFile, "c", "", "(deprecated alias of -config)")
-	fs.StringVar(&cfg.ConfigFile, "config", "tidb-lightning.toml", "tidb-lightning configuration file")
-	fs.BoolVar(&cfg.DoCompact, "compact", false, "do manual compaction on the target cluster, run then exit")
-	fs.StringVar(&cfg.SwitchMode, "switch-mode", "", "switch tikv into import mode or normal mode, values can be ['import', 'normal'], run then exit")
-	fs.BoolVar(&cfg.printVersion, "V", false, "print version of lightning")
+	fs.StringVar(&cfg.ConfigFile, "config", "", "tidb-lightning configuration file")
+	printVersion := fs.Bool("V", false, "print version of lightning")
+
+	logLevel := fs.String("L", "", `log level: info, debug, warn, error, fatal (default "info")`)
+	logFilePath := fs.String("log-file", "", "log file path")
+	tidbHost := fs.String("tidb-host", "", "TiDB server host")
+	tidbPort := fs.Int("tidb-port", 0, "TiDB server port (default 4000)")
+	tidbUser := fs.String("tidb-user", "", "TiDB user name to connect")
+	tidbStatusPort := fs.Int("tidb-status", 0, "TiDB server status port (default 10080)")
+	pdAddr := fs.String("pd-urls", "", "PD endpoint address")
+	dataSrcPath := fs.String("d", "", "Directory of the dump to import")
+	importerAddr := fs.String("importer", "", "address (host:port) to connect to tikv-importer")
 
 	if err := fs.Parse(args); err != nil {
 		return nil, errors.Trace(err)
 	}
+	if *printVersion {
+		fmt.Println(common.GetRawInfo())
+		return nil, flag.ErrHelp
+	}
 
 	if err := cfg.Load(); err != nil {
 		return nil, errors.Trace(err)
 	}
+
+	if *logLevel != "" {
+		cfg.App.LogConfig.Level = *logLevel
+	}
+	if *logFilePath != "" {
+		cfg.App.LogConfig.File = *logFilePath
+	}
+	if *tidbHost != "" {
+		cfg.TiDB.Host = *tidbHost
+	}
+	if *tidbPort != 0 {
+		cfg.TiDB.Port = *tidbPort
+	}
+	if *tidbStatusPort != 0 {
+		cfg.TiDB.StatusPort = *tidbStatusPort
+	}
+	if *tidbUser != "" {
+		cfg.TiDB.User = *tidbUser
+	}
+	if *pdAddr != "" {
+		cfg.TiDB.PdAddr = *pdAddr
+	}
+	if *dataSrcPath != "" {
+		cfg.Mydumper.SourceDir = *dataSrcPath
+	}
+	if *importerAddr != "" {
+		cfg.TikvImporter.Addr = *importerAddr
+	}
+
+	cfg.Adjust()
+
 	return cfg, nil
 }
 
 func (cfg *Config) Load() error {
-	if cfg.printVersion {
-		fmt.Println(common.GetRawInfo())
-		return flag.ErrHelp
+	// use standard config if unspecified.
+	if cfg.ConfigFile == "" {
+		return nil
 	}
 
 	data, err := ioutil.ReadFile(cfg.ConfigFile)
@@ -234,6 +282,10 @@ func (cfg *Config) Load() error {
 		return errors.Annotate(err, "invalid config: `mydumper.tidb.sql_mode` must be a valid SQL_MODE")
 	}
 
+	return nil
+}
+
+func (cfg *Config) Adjust() {
 	// handle mydumper
 	if cfg.Mydumper.BatchSize <= 0 {
 		cfg.Mydumper.BatchSize = 100 * _G
@@ -262,5 +314,4 @@ func (cfg *Config) Load() error {
 			cfg.Checkpoint.DSN = "/tmp/" + cfg.Checkpoint.Schema + ".pb"
 		}
 	}
-	return nil
 }

lightning/lightning.go

@@ -21,12 +21,10 @@ import (
 	"sync"
 
 	"github.com/pingcap/errors"
-	sstpb "github.com/pingcap/kvproto/pkg/import_sstpb"
 	"github.com/prometheus/client_golang/prometheus/promhttp"
 
 	"github.com/pingcap/tidb-lightning/lightning/common"
 	"github.com/pingcap/tidb-lightning/lightning/config"
-	"github.com/pingcap/tidb-lightning/lightning/kv"
 	"github.com/pingcap/tidb-lightning/lightning/mydump"
 	"github.com/pingcap/tidb-lightning/lightning/restore"
 )
@@ -72,10 +70,6 @@ func (l *Lightning) Run() error {
 		common.AppLogger.Infof("cfg %s", l.cfg)
 	})
 
-	if l.handleCommandFlagsAndExits() {
-		return nil
-	}
-
 	l.wg.Add(1)
 	var err error
 	go func() {
@@ -86,34 +80,6 @@ func (l *Lightning) Run() error {
 	return errors.Trace(err)
 }
 
-func (l *Lightning) handleCommandFlagsAndExits() (exits bool) {
-	if l.cfg.DoCompact {
-		err := l.doCompact()
-		if err != nil {
-			common.AppLogger.Fatalf("compact error %s", errors.ErrorStack(err))
-		}
-		return true
-	}
-
-	if mode := l.cfg.SwitchMode; mode != "" {
-		var err error
-		switch mode {
-		case config.ImportMode:
-			err = l.switchMode(sstpb.SwitchMode_Import)
-		case config.NormalMode:
-			err = l.switchMode(sstpb.SwitchMode_Normal)
-		default:
-			common.AppLogger.Fatalf("invalid mode %s, must use %s or %s", mode, config.ImportMode, config.NormalMode)
-		}
-		if err != nil {
-			common.AppLogger.Fatalf("switch mode error %v", errors.ErrorStack(err))
-		}
-		common.AppLogger.Infof("switch mode to %s", mode)
-		return true
-	}
-	return false
-}
-
 func (l *Lightning) run() error {
 	mdl, err := mydump.NewMyDumpLoader(l.cfg)
 	if err != nil {
@@ -134,38 +100,6 @@ func (l *Lightning) run() error {
 	return errors.Trace(err)
 }
 
-func (l *Lightning) doCompact() error {
-	ctx := context.Background()
-
-	importer, err := kv.NewImporter(ctx, l.cfg.TikvImporter.Addr, l.cfg.TiDB.PdAddr)
-	if err != nil {
-		return errors.Trace(err)
-	}
-	defer importer.Close()
-
-	if err := importer.Compact(ctx, restore.FullLevelCompact); err != nil {
-		return errors.Trace(err)
-	}
-
-	return nil
-}
-
-func (l *Lightning) switchMode(mode sstpb.SwitchMode) error {
-	ctx := context.Background()
-
-	importer, err := kv.NewImporter(ctx, l.cfg.TikvImporter.Addr, l.cfg.TiDB.PdAddr)
-	if err != nil {
-		return errors.Trace(err)
-	}
-	defer importer.Close()
-
-	if err := importer.SwitchMode(ctx, mode); err != nil {
-		return errors.Trace(err)
-	}
-
-	return nil
-}
-
 func (l *Lightning) Stop() {
 	l.shutdown()
 	l.wg.Wait()

tests/_utils/run_lightning

@@ -15,6 +15,8 @@
 
 set -eu
 TEST_DIR=/tmp/lightning_test_result
+CONFIG="${1-config}"
+shift
 
-echo "[$(date)] <<<<<< RUNNING TEST FOR: tests/$TEST_NAME/${1-config}.toml >>>>>>" >> "$TEST_DIR/lightning.log"
-bin/tidb-lightning.test -test.coverprofile="$TEST_DIR/cov.$TEST_NAME.$$.out" DEVEL -config "tests/$TEST_NAME/${1-config}.toml"
+echo "[$(date)] <<<<<< RUNNING TEST FOR: tests/$TEST_NAME/$CONFIG.toml >>>>>>" >> "$TEST_DIR/lightning.log"
+bin/tidb-lightning.test -test.coverprofile="$TEST_DIR/cov.$TEST_NAME.$$.out" DEVEL -config "tests/$TEST_NAME/$CONFIG.toml" "$@"

tests/cmdline_override/config.toml

@@ -0,0 +1,18 @@
+[lightning]
+check-requirements = false
+file = "/tmp/xyzxyzxyz"
+level = "xyzxyzxyz"
+
+[tikv-importer]
+addr = "xyzxyzxyz"
+
+[mydumper]
+data-source-dir = "xyzxyzxyz"
+
+[tidb]
+host = "xyzxyzxyz"
+port = 12345678
+user = "xyzxyzxyz"
+status-port = 12345678
+pd-addr = "xyzxyzxyz"
+

tests/cmdline_override/data/cmdline_override-schema-create.sql

@@ -0,0 +1 @@
+create database cmdline_override;

tests/cmdline_override/data/cmdline_override.t-schema.sql

@@ -0,0 +1 @@
+create table t (a int);

tests/cmdline_override/data/cmdline_override.t.sql

@@ -0,0 +1 @@
+insert into t values (15);

tests/cmdline_override/run.sh

@@ -0,0 +1,17 @@
+#!/bin/sh
+
+set -eux
+
+run_lightning config \
+    -L info \
+    --log-file "$TEST_DIR/lightning.log" \
+    --tidb-host 127.0.0.1 \
+    --tidb-port 4000 \
+    --tidb-user root \
+    --tidb-status 10080 \
+    --pd-urls 127.0.0.1:2379 \
+    -d tests/cmdline_override/data \
+    --importer 127.0.0.1:8808
+
+run_sql 'SELECT * FROM cmdline_override.t'
+check_contains 'a: 15'