1// Copyright (C) MongoDB, Inc. 2014-present.
2//
3// Licensed under the Apache License, Version 2.0 (the "License"); you may
4// not use this file except in compliance with the License. You may obtain
5// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
6
7// Main package for the mongostat tool.
8package main
9
10import (
11	"os"
12	"strconv"
13	"strings"
14	"time"
15
16	"github.com/mongodb/mongo-tools/common/log"
17	"github.com/mongodb/mongo-tools/common/options"
18	"github.com/mongodb/mongo-tools/common/password"
19	"github.com/mongodb/mongo-tools/common/signals"
20	"github.com/mongodb/mongo-tools/common/util"
21	"github.com/mongodb/mongo-tools/mongostat"
22	"github.com/mongodb/mongo-tools/mongostat/stat_consumer"
23	"github.com/mongodb/mongo-tools/mongostat/stat_consumer/line"
24	"github.com/mongodb/mongo-tools/mongostat/status"
25)
26
27// optionKeyNames interprets the CLI options Columns and AppendColumns into
28// the internal keyName mapping.
29func optionKeyNames(option string) map[string]string {
30	kn := make(map[string]string)
31	columns := strings.Split(option, ",")
32	for _, column := range columns {
33		naming := strings.Split(column, "=")
34		if len(naming) == 1 {
35			kn[naming[0]] = naming[0]
36		} else {
37			kn[naming[0]] = naming[1]
38		}
39	}
40	return kn
41}
42
43// optionCustomHeaders interprets the CLI options Columns and AppendColumns
44// into a list of custom headers.
45func optionCustomHeaders(option string) (headers []string) {
46	columns := strings.Split(option, ",")
47	for _, column := range columns {
48		naming := strings.Split(column, "=")
49		headers = append(headers, naming[0])
50	}
51	return
52}
53
54func main() {
55	// initialize command-line opts
56	opts := options.New(
57		"mongostat",
58		mongostat.Usage,
59		options.EnabledOptions{Connection: true, Auth: true, Namespace: false, URI: true})
60	opts.UseReadOnlyHostDescription()
61
62	// add mongostat-specific options
63	statOpts := &mongostat.StatOptions{}
64	opts.AddOptions(statOpts)
65
66	interactiveOption := opts.FindOptionByLongName("interactive")
67	if _, available := stat_consumer.FormatterConstructors["interactive"]; !available {
68		// make --interactive inaccessible
69		interactiveOption.LongName = ""
70		interactiveOption.ShortName = 0
71	}
72
73	args, err := opts.ParseArgs(os.Args[1:])
74	if err != nil {
75		log.Logvf(log.Always, "error parsing command line options: %v", err)
76		log.Logvf(log.Always, "try 'mongostat --help' for more information")
77		os.Exit(util.ExitBadOptions)
78	}
79
80	log.SetVerbosity(opts.Verbosity)
81	signals.Handle()
82
83	sleepInterval := 1
84	if len(args) > 0 {
85		if len(args) != 1 {
86			log.Logvf(log.Always, "too many positional arguments: %v", args)
87			log.Logvf(log.Always, "try 'mongostat --help' for more information")
88			os.Exit(util.ExitBadOptions)
89		}
90		sleepInterval, err = strconv.Atoi(args[0])
91		if err != nil {
92			log.Logvf(log.Always, "invalid sleep interval: %v", args[0])
93			os.Exit(util.ExitBadOptions)
94		}
95		if sleepInterval < 1 {
96			log.Logvf(log.Always, "sleep interval must be at least 1 second")
97			os.Exit(util.ExitBadOptions)
98		}
99	}
100
101	// print help, if specified
102	if opts.PrintHelp(false) {
103		return
104	}
105
106	// print version, if specified
107	if opts.PrintVersion() {
108		return
109	}
110
111	// verify uri options and log them
112	opts.URI.LogUnsupportedOptions()
113
114	if opts.Auth.Username != "" && opts.GetAuthenticationDatabase() == "" && !opts.Auth.RequiresExternalDB() {
115		// add logic to have different error if using uri
116		if opts.URI != nil && opts.URI.ConnectionString != "" {
117			log.Logvf(log.Always, "authSource is required when authenticating against a non $external database")
118			os.Exit(util.ExitBadOptions)
119		}
120
121		log.Logvf(log.Always, "--authenticationDatabase is required when authenticating against a non $external database")
122		os.Exit(util.ExitBadOptions)
123	}
124
125	if statOpts.Interactive && statOpts.Json {
126		log.Logvf(log.Always, "cannot use output formats --json and --interactive together")
127		os.Exit(util.ExitBadOptions)
128	}
129
130	if statOpts.Deprecated && !statOpts.Json {
131		log.Logvf(log.Always, "--useDeprecatedJsonKeys can only be used when --json is also specified")
132		os.Exit(util.ExitBadOptions)
133	}
134
135	if statOpts.Columns != "" && statOpts.AppendColumns != "" {
136		log.Logvf(log.Always, "-O cannot be used if -o is also specified")
137		os.Exit(util.ExitBadOptions)
138	}
139
140	if statOpts.HumanReadable != "true" && statOpts.HumanReadable != "false" {
141		log.Logvf(log.Always, "--humanReadable must be set to either 'true' or 'false'")
142		os.Exit(util.ExitBadOptions)
143	}
144
145	// we have to check this here, otherwise the user will be prompted
146	// for a password for each discovered node
147	if opts.Auth.ShouldAskForPassword() {
148		opts.Auth.Password = password.Prompt()
149	}
150
151	var factory stat_consumer.FormatterConstructor
152	if statOpts.Json {
153		factory = stat_consumer.FormatterConstructors["json"]
154	} else if statOpts.Interactive {
155		factory = stat_consumer.FormatterConstructors["interactive"]
156	} else {
157		factory = stat_consumer.FormatterConstructors[""]
158	}
159	formatter := factory(statOpts.RowCount, !statOpts.NoHeaders)
160
161	cliFlags := 0
162	if statOpts.Columns == "" {
163		cliFlags = line.FlagAlways
164		if statOpts.Discover {
165			cliFlags |= line.FlagDiscover
166			cliFlags |= line.FlagHosts
167		}
168		if statOpts.All {
169			cliFlags |= line.FlagAll
170		}
171		if strings.Contains(opts.Host, ",") {
172			cliFlags |= line.FlagHosts
173		}
174	}
175
176	var customHeaders []string
177	if statOpts.Columns != "" {
178		customHeaders = optionCustomHeaders(statOpts.Columns)
179	} else if statOpts.AppendColumns != "" {
180		customHeaders = optionCustomHeaders(statOpts.AppendColumns)
181	}
182
183	var keyNames map[string]string
184	if statOpts.Deprecated {
185		keyNames = line.DeprecatedKeyMap()
186	} else if statOpts.Columns == "" {
187		keyNames = line.DefaultKeyMap()
188	} else {
189		keyNames = optionKeyNames(statOpts.Columns)
190	}
191	if statOpts.AppendColumns != "" {
192		addKN := optionKeyNames(statOpts.AppendColumns)
193		for k, v := range addKN {
194			keyNames[k] = v
195		}
196	}
197
198	readerConfig := &status.ReaderConfig{
199		HumanReadable: statOpts.HumanReadable == "true",
200	}
201	if statOpts.Json {
202		readerConfig.TimeFormat = "15:04:05"
203	}
204
205	consumer := stat_consumer.NewStatConsumer(cliFlags, customHeaders,
206		keyNames, readerConfig, formatter, os.Stdout)
207	seedHosts := util.CreateConnectionAddrs(opts.Host, opts.Port)
208	var cluster mongostat.ClusterMonitor
209	if statOpts.Discover || len(seedHosts) > 1 {
210		cluster = &mongostat.AsyncClusterMonitor{
211			ReportChan:    make(chan *status.ServerStatus),
212			ErrorChan:     make(chan *status.NodeError),
213			LastStatLines: map[string]*line.StatLine{},
214			Consumer:      consumer,
215		}
216	} else {
217		cluster = &mongostat.SyncClusterMonitor{
218			ReportChan: make(chan *status.ServerStatus),
219			ErrorChan:  make(chan *status.NodeError),
220			Consumer:   consumer,
221		}
222	}
223
224	var discoverChan chan string
225	if statOpts.Discover {
226		discoverChan = make(chan string, 128)
227	}
228
229	opts.Direct = true
230	stat := &mongostat.MongoStat{
231		Options:       opts,
232		StatOptions:   statOpts,
233		Nodes:         map[string]*mongostat.NodeMonitor{},
234		Discovered:    discoverChan,
235		SleepInterval: time.Duration(sleepInterval) * time.Second,
236		Cluster:       cluster,
237	}
238
239	for _, v := range seedHosts {
240		stat.AddNewNode(v)
241	}
242
243	// kick it off
244	err = stat.Run()
245	formatter.Finish()
246	if err != nil {
247		log.Logvf(log.Always, "Failed: %v", err)
248		os.Exit(util.ExitError)
249	}
250}
251