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
54var (
55	VersionStr = "built-without-version-string"
56	GitCommit  = "build-without-git-commit"
57)
58
59func main() {
60	// initialize command-line opts
61	opts := options.New(
62		"mongostat", VersionStr, GitCommit,
63		mongostat.Usage,
64		options.EnabledOptions{Connection: true, Auth: true, Namespace: false, URI: true})
65	opts.UseReadOnlyHostDescription()
66
67	// add mongostat-specific options
68	statOpts := &mongostat.StatOptions{}
69	opts.AddOptions(statOpts)
70
71	interactiveOption := opts.FindOptionByLongName("interactive")
72	if _, available := stat_consumer.FormatterConstructors["interactive"]; !available {
73		// make --interactive inaccessible
74		interactiveOption.LongName = ""
75		interactiveOption.ShortName = 0
76	}
77
78	args, err := opts.ParseArgs(os.Args[1:])
79	if err != nil {
80		log.Logvf(log.Always, "error parsing command line options: %v", err)
81		log.Logvf(log.Always, util.ShortUsage("mongostat"))
82		os.Exit(util.ExitFailure)
83	}
84
85	log.SetVerbosity(opts.Verbosity)
86	signals.Handle()
87
88	sleepInterval := 1
89	if len(args) > 0 {
90		if len(args) != 1 {
91			log.Logvf(log.Always, "too many positional arguments: %v", args)
92			log.Logvf(log.Always, util.ShortUsage("mongostat"))
93			os.Exit(util.ExitFailure)
94		}
95		sleepInterval, err = strconv.Atoi(args[0])
96		if err != nil {
97			log.Logvf(log.Always, "invalid sleep interval: %v", args[0])
98			os.Exit(util.ExitFailure)
99		}
100		if sleepInterval < 1 {
101			log.Logvf(log.Always, "sleep interval must be at least 1 second")
102			os.Exit(util.ExitFailure)
103		}
104	}
105
106	// print help, if specified
107	if opts.PrintHelp(false) {
108		return
109	}
110
111	// print version, if specified
112	if opts.PrintVersion() {
113		return
114	}
115
116	// verify uri options and log them
117	opts.URI.LogUnsupportedOptions()
118
119	if opts.Auth.Username != "" && opts.GetAuthenticationDatabase() == "" && !opts.Auth.RequiresExternalDB() {
120		// add logic to have different error if using uri
121		if opts.URI != nil && opts.URI.ConnectionString != "" {
122			log.Logvf(log.Always, "authSource is required when authenticating against a non $external database")
123			os.Exit(util.ExitFailure)
124		}
125
126		log.Logvf(log.Always, "--authenticationDatabase is required when authenticating against a non $external database")
127		os.Exit(util.ExitFailure)
128	}
129
130	if statOpts.Interactive && statOpts.Json {
131		log.Logvf(log.Always, "cannot use output formats --json and --interactive together")
132		os.Exit(util.ExitFailure)
133	}
134
135	if statOpts.Deprecated && !statOpts.Json {
136		log.Logvf(log.Always, "--useDeprecatedJsonKeys can only be used when --json is also specified")
137		os.Exit(util.ExitFailure)
138	}
139
140	if statOpts.Columns != "" && statOpts.AppendColumns != "" {
141		log.Logvf(log.Always, "-O cannot be used if -o is also specified")
142		os.Exit(util.ExitFailure)
143	}
144
145	if statOpts.HumanReadable != "true" && statOpts.HumanReadable != "false" {
146		log.Logvf(log.Always, "--humanReadable must be set to either 'true' or 'false'")
147		os.Exit(util.ExitFailure)
148	}
149
150	// we have to check this here, otherwise the user will be prompted
151	// for a password for each discovered node
152	if opts.Auth.ShouldAskForPassword() {
153		pass, err := password.Prompt()
154		if err != nil {
155			log.Logvf(log.Always, "Failed: %v", err)
156			os.Exit(util.ExitFailure)
157		}
158		opts.Auth.Password = pass
159	}
160
161	var factory stat_consumer.FormatterConstructor
162	if statOpts.Json {
163		factory = stat_consumer.FormatterConstructors["json"]
164	} else if statOpts.Interactive {
165		factory = stat_consumer.FormatterConstructors["interactive"]
166	} else {
167		factory = stat_consumer.FormatterConstructors[""]
168	}
169	formatter := factory(statOpts.RowCount, !statOpts.NoHeaders)
170
171	cliFlags := 0
172	if statOpts.Columns == "" {
173		cliFlags = line.FlagAlways
174		if statOpts.Discover {
175			cliFlags |= line.FlagDiscover
176			cliFlags |= line.FlagHosts
177		}
178		if statOpts.All {
179			cliFlags |= line.FlagAll
180		}
181		if strings.Contains(opts.Host, ",") {
182			cliFlags |= line.FlagHosts
183		}
184	}
185
186	var customHeaders []string
187	if statOpts.Columns != "" {
188		customHeaders = optionCustomHeaders(statOpts.Columns)
189	} else if statOpts.AppendColumns != "" {
190		customHeaders = optionCustomHeaders(statOpts.AppendColumns)
191	}
192
193	var keyNames map[string]string
194	if statOpts.Deprecated {
195		keyNames = line.DeprecatedKeyMap()
196	} else if statOpts.Columns == "" {
197		keyNames = line.DefaultKeyMap()
198	} else {
199		keyNames = optionKeyNames(statOpts.Columns)
200	}
201	if statOpts.AppendColumns != "" {
202		addKN := optionKeyNames(statOpts.AppendColumns)
203		for k, v := range addKN {
204			keyNames[k] = v
205		}
206	}
207
208	readerConfig := &status.ReaderConfig{
209		HumanReadable: statOpts.HumanReadable == "true",
210	}
211	if statOpts.Json {
212		readerConfig.TimeFormat = "15:04:05"
213	}
214
215	consumer := stat_consumer.NewStatConsumer(cliFlags, customHeaders,
216		keyNames, readerConfig, formatter, os.Stdout)
217	seedHosts := util.CreateConnectionAddrs(opts.Host, opts.Port)
218	var cluster mongostat.ClusterMonitor
219	if statOpts.Discover || len(seedHosts) > 1 {
220		cluster = &mongostat.AsyncClusterMonitor{
221			ReportChan:    make(chan *status.ServerStatus),
222			ErrorChan:     make(chan *status.NodeError),
223			LastStatLines: map[string]*line.StatLine{},
224			Consumer:      consumer,
225		}
226	} else {
227		cluster = &mongostat.SyncClusterMonitor{
228			ReportChan: make(chan *status.ServerStatus),
229			ErrorChan:  make(chan *status.NodeError),
230			Consumer:   consumer,
231		}
232	}
233
234	var discoverChan chan string
235	if statOpts.Discover {
236		discoverChan = make(chan string, 128)
237	}
238
239	opts.Direct = true
240	stat := &mongostat.MongoStat{
241		Options:       opts,
242		StatOptions:   statOpts,
243		Nodes:         map[string]*mongostat.NodeMonitor{},
244		Discovered:    discoverChan,
245		SleepInterval: time.Duration(sleepInterval) * time.Second,
246		Cluster:       cluster,
247	}
248
249	for _, v := range seedHosts {
250		if err := stat.AddNewNode(v); err != nil {
251			log.Logv(log.Always, err.Error())
252			os.Exit(util.ExitFailure)
253		}
254	}
255
256	// kick it off
257	err = stat.Run()
258	for _, monitor := range stat.Nodes {
259		monitor.Disconnect()
260	}
261	formatter.Finish()
262	if err != nil {
263		log.Logvf(log.Always, "Failed: %v", err)
264		os.Exit(util.ExitFailure)
265	}
266}
267