github.com/rohankumardubey/aresdb@v0.0.2-0.20190517170215-e54e3ca06b9c/cmd/subscriber/main.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"github.com/uber/aresdb/subscriber/common/consumer/kafka"
     7  	"github.com/uber/aresdb/subscriber/common/sink"
     8  	"os"
     9  	"runtime"
    10  
    11  	"github.com/spf13/cobra"
    12  	"github.com/uber-go/tally"
    13  	"github.com/uber/aresdb/subscriber/common/job"
    14  	"github.com/uber/aresdb/subscriber/common/message"
    15  	"github.com/uber/aresdb/subscriber/common/rules"
    16  	"github.com/uber/aresdb/subscriber/config"
    17  	"github.com/uber/aresdb/utils"
    18  	cfgfx "go.uber.org/config"
    19  	"go.uber.org/fx"
    20  	"go.uber.org/zap"
    21  )
    22  
    23  type Params struct {
    24  	fx.In
    25  
    26  	Lifecycle fx.Lifecycle
    27  }
    28  
    29  type Result struct {
    30  	fx.Out
    31  
    32  	Environment utils.EnvironmentContext
    33  	Provider    cfgfx.Provider
    34  	Logger      *zap.Logger
    35  	Scope       tally.Scope
    36  	Consumer    job.NewConsumer
    37  	Decoder     job.NewDecoder
    38  	Sink        job.NewSink
    39  }
    40  
    41  func main() {
    42  	module := fx.Provide(Init)
    43  	Execute(module, config.Module, rules.Module, job.Module)
    44  }
    45  
    46  func Execute(opts ...fx.Option) {
    47  	cmd := &cobra.Command{
    48  		Use:     "ares-subscriber",
    49  		Short:   "AresDB subscriber",
    50  		Long:    `AresDB subscriber ingest data into AresDB`,
    51  		Example: `ares-subscriber -c config/staging/staging.yaml -r config -j summary -a summaryProd`,
    52  		Run: func(cmd *cobra.Command, args []string) {
    53  			runtime.GOMAXPROCS(runtime.NumCPU())
    54  			fx.New(opts...).Run()
    55  		},
    56  	}
    57  	addFlags(cmd)
    58  	cmd.Execute()
    59  }
    60  
    61  func addFlags(cmd *cobra.Command) {
    62  	cmd.Flags().StringVar(&config.ConfigRootPath, "r", "subscriber/config", "Config root path")
    63  	cmd.Flags().StringVar(&config.ConfigFile, "c", "subscriber/config/staging/staging.yaml", "Config file path")
    64  	cmd.Flags().StringVar(&config.ActiveJobNameSpace, "j", "jobNameSpace", "Job namespace")
    65  	cmd.Flags().StringVar(&config.ActiveAresNameSpace, "a", "aresNameSpace", "Ares cluster namespace")
    66  
    67  	if cmd.MarkFlagRequired("j") != nil {
    68  		panic(`missing "-j jobNameSpace". Job namespace is required`)
    69  	}
    70  
    71  	if cmd.MarkFlagRequired("a") != nil {
    72  		panic(`missing "-a aresNameSpace". Ares cluster namespace is required`)
    73  	}
    74  
    75  	fmt.Printf("jobNameSpace=%s, aresNameSpace=%s\n", config.ActiveJobNameSpace, config.ActiveAresNameSpace)
    76  }
    77  
    78  func newDefaultEnvironmentCtx() utils.EnvironmentContext {
    79  	return utils.EnvironmentContext{
    80  		Environment:        os.Getenv(utils.EnvironmentKey),
    81  		RuntimeEnvironment: os.Getenv(utils.RuntimeEnvironmentKey),
    82  		Zone:               os.Getenv(utils.ZoneKey),
    83  		Hostname:           os.Getenv(utils.HostnameKey),
    84  		Deployment:         os.Getenv(utils.DeploymentKey),
    85  		SystemPort:         os.Getenv(utils.PortSystemKey),
    86  		ApplicationID:      os.Getenv(utils.AppIDKey),
    87  		InstanceID:         os.Getenv(utils.InstanceIDKey),
    88  	}
    89  }
    90  
    91  func newDefaultConfig() cfgfx.Provider {
    92  	if config.ConfigFile == "" {
    93  		panic(`missing "-c configFile"`)
    94  	}
    95  
    96  	provider, err := cfgfx.NewYAMLProviderFromFiles(config.ConfigFile)
    97  	if err != nil {
    98  		panic(utils.StackError(err, fmt.Sprintf("Failed to load config %s", config.ConfigFile)))
    99  	}
   100  	return provider
   101  }
   102  
   103  func newDefaultLogger(params Params, env utils.EnvironmentContext) *zap.Logger {
   104  	logger := zap.NewExample()
   105  	undoGlobals := zap.ReplaceGlobals(logger)
   106  	undoHijack := zap.RedirectStdLog(logger)
   107  	params.Lifecycle.Append(fx.Hook{
   108  		OnStart: func(context.Context) error {
   109  			logger.Info("starting up with environment", zap.Any("environment", env.Environment))
   110  			return nil
   111  		},
   112  		OnStop: func(context.Context) error {
   113  			undoHijack()
   114  			undoGlobals()
   115  			logger.Sync()
   116  			return nil
   117  		},
   118  	})
   119  	return logger
   120  }
   121  
   122  func newDefaultScope() tally.Scope {
   123  	return tally.NewTestScope("test", nil)
   124  }
   125  
   126  func Init(params Params) Result {
   127  	env := newDefaultEnvironmentCtx()
   128  	provider := newDefaultConfig()
   129  	logger := newDefaultLogger(params, env)
   130  	scope := newDefaultScope()
   131  
   132  	return Result{
   133  		Environment: env,
   134  		Provider:    provider,
   135  		Logger:      logger,
   136  		Scope:       scope,
   137  		Consumer:    kafka.NewKafkaConsumer,
   138  		Decoder:     message.NewDefaultDecoder,
   139  		Sink:        sink.NewAresDatabase,
   140  	}
   141  }