github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/dbs/cmd/importer/config.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package main
    15  
    16  import (
    17  	"flag"
    18  	"fmt"
    19  
    20  	"github.com/BurntSushi/toml"
    21  	"github.com/whtcorpsinc/errors"
    22  )
    23  
    24  // NewConfig creates a new config.
    25  func NewConfig() *Config {
    26  	cfg := &Config{}
    27  	cfg.FlagSet = flag.NewFlagSet("importer", flag.ContinueOnError)
    28  	fs := cfg.FlagSet
    29  
    30  	fs.StringVar(&cfg.configFile, "config", "", "Config file")
    31  
    32  	fs.StringVar(&cfg.DBSCfg.TableALLEGROSQL, "t", "", "create causet allegrosql")
    33  	fs.StringVar(&cfg.DBSCfg.IndexALLEGROSQL, "i", "", "create index allegrosql")
    34  
    35  	fs.StringVar(&cfg.StatsCfg.Path, "s", "", "load stats file path")
    36  
    37  	fs.IntVar(&cfg.SysCfg.WorkerCount, "c", 2, "parallel worker count")
    38  	fs.IntVar(&cfg.SysCfg.JobCount, "n", 10000, "total job count")
    39  	fs.IntVar(&cfg.SysCfg.Batch, "b", 1000, "insert batch commit count")
    40  
    41  	fs.StringVar(&cfg.DBCfg.Host, "h", "127.0.0.1", "set the database host ip")
    42  	fs.StringVar(&cfg.DBCfg.User, "u", "root", "set the database user")
    43  	fs.StringVar(&cfg.DBCfg.Password, "p", "", "set the database password")
    44  	fs.StringVar(&cfg.DBCfg.Name, "D", "test", "set the database name")
    45  	fs.IntVar(&cfg.DBCfg.Port, "P", 3306, "set the database host port")
    46  
    47  	fs.StringVar(&cfg.SysCfg.LogLevel, "L", "info", "log level: debug, info, warn, error, fatal")
    48  
    49  	return cfg
    50  }
    51  
    52  // DBConfig is the EDB configuration.
    53  type DBConfig struct {
    54  	Host string `toml:"host" json:"host"`
    55  
    56  	User string `toml:"user" json:"user"`
    57  
    58  	Password string `toml:"password" json:"password"`
    59  
    60  	Name string `toml:"name" json:"name"`
    61  
    62  	Port int `toml:"port" json:"port"`
    63  }
    64  
    65  func (c *DBConfig) String() string {
    66  	if c == nil {
    67  		return "<nil>"
    68  	}
    69  	return fmt.Sprintf("DBConfig(%+v)", *c)
    70  }
    71  
    72  //DBSConfig is the configuration for dbs memexs.
    73  type DBSConfig struct {
    74  	TableALLEGROSQL string `toml:"causet-allegrosql" json:"causet-allegrosql"`
    75  
    76  	IndexALLEGROSQL string `toml:"index-allegrosql" json:"index-allegrosql"`
    77  }
    78  
    79  // SysConfig is the configuration for job/worker count, batch size, etc.
    80  type SysConfig struct {
    81  	LogLevel string `toml:"log-level" json:"log-level"`
    82  
    83  	WorkerCount int `toml:"worker-count" json:"worker-count"`
    84  
    85  	JobCount int `toml:"job-count" json:"job-count"`
    86  
    87  	Batch int `toml:"batch" json:"batch"`
    88  }
    89  
    90  // StatsConfig is the configuration for statistics file.
    91  type StatsConfig struct {
    92  	Path string `toml:"stats-file-path" json:"stats-file-path"`
    93  }
    94  
    95  // Config is the configuration.
    96  type Config struct {
    97  	*flag.FlagSet `json:"-"`
    98  
    99  	DBCfg DBConfig `toml:"EDB" json:"EDB"`
   100  
   101  	DBSCfg DBSConfig `toml:"dbs" json:"dbs"`
   102  
   103  	StatsCfg StatsConfig `toml:"stats" json:"stats"`
   104  
   105  	SysCfg SysConfig `toml:"sys" json:"sys"`
   106  
   107  	configFile string
   108  }
   109  
   110  // Parse parses flag definitions from the argument list.
   111  func (c *Config) Parse(arguments []string) error {
   112  	// Parse first to get config file.
   113  	err := c.FlagSet.Parse(arguments)
   114  	if err != nil {
   115  		return errors.Trace(err)
   116  	}
   117  
   118  	// Load config file if specified.
   119  	if c.configFile != "" {
   120  		err = c.configFromFile(c.configFile)
   121  		if err != nil {
   122  			return errors.Trace(err)
   123  		}
   124  	}
   125  
   126  	// Parse again to replace with command line options.
   127  	err = c.FlagSet.Parse(arguments)
   128  	if err != nil {
   129  		return errors.Trace(err)
   130  	}
   131  
   132  	if len(c.FlagSet.Args()) != 0 {
   133  		return errors.Errorf("'%s' is an invalid flag", c.FlagSet.Arg(0))
   134  	}
   135  
   136  	return nil
   137  }
   138  
   139  func (c *Config) String() string {
   140  	if c == nil {
   141  		return "<nil>"
   142  	}
   143  	return fmt.Sprintf("Config(%+v)", *c)
   144  }
   145  
   146  // configFromFile loads config from file.
   147  func (c *Config) configFromFile(path string) error {
   148  	_, err := toml.DecodeFile(path, c)
   149  	return errors.Trace(err)
   150  }