chore(outputs.mongodb): Cleanup plugin

Author: srebhan

plugins/outputs/mongodb/mongodb.go

@@ -25,32 +25,6 @@ import (
 //go:embed sample.conf
 var sampleConfig string
 
-func (s *MongoDB) getCollections(ctx context.Context) error {
-	collections, err := s.client.Database(s.MetricDatabase).ListCollections(ctx, bson.M{})
-	if err != nil {
-		return fmt.Errorf("unable to execute ListCollections: %w", err)
-	}
-	s.collections = make(map[string]bson.M, collections.RemainingBatchLength())
-	for collections.Next(ctx) {
-		var collection bson.M
-		if err = collections.Decode(&collection); err != nil {
-			return fmt.Errorf("unable to decode ListCollections: %w", err)
-		}
-		name, ok := collection["name"].(string)
-		if !ok {
-			return fmt.Errorf("non-string name in %v", collection)
-		}
-		s.collections[name] = collection
-	}
-	return nil
-}
-
-func (s *MongoDB) insertDocument(ctx context.Context, databaseCollection string, bdoc bson.D) error {
-	collection := s.client.Database(s.MetricDatabase).Collection(databaseCollection)
-	_, err := collection.InsertOne(ctx, &bdoc)
-	return err
-}
-
 type MongoDB struct {
 	Dsn                 string          `toml:"dsn"`
 	AuthenticationType  string          `toml:"authentication"`
@@ -61,10 +35,11 @@ type MongoDB struct {
 	ServerSelectTimeout config.Duration `toml:"timeout"`
 	TTL                 config.Duration `toml:"ttl"`
 	Log                 telegraf.Logger `toml:"-"`
-	client              *mongo.Client
-	clientOptions       *options.ClientOptions
-	collections         map[string]bson.M
 	tls.ClientConfig
+
+	client      *mongo.Client
+	options     *options.ClientOptions
+	collections map[string]bool
 }
 
 func (*MongoDB) SampleConfig() string {
@@ -92,7 +67,7 @@ func (s *MongoDB) Init() error {
 	}
 
 	serverAPIOptions := options.ServerAPI(options.ServerAPIVersion1) // use new mongodb versioned api
-	s.clientOptions = options.Client().SetServerAPIOptions(serverAPIOptions)
+	s.options = options.Client().SetServerAPIOptions(serverAPIOptions)
 
 	switch s.AuthenticationType {
 	case "", "NONE":
@@ -120,7 +95,7 @@ func (s *MongoDB) Init() error {
 		}
 		username.Destroy()
 		password.Destroy()
-		s.clientOptions.SetAuth(credential)
+		s.options.SetAuth(credential)
 	case "PLAIN":
 		if s.Username.Empty() {
 			return errors.New("authentication for PLAIN must specify a username")
@@ -148,7 +123,7 @@ func (s *MongoDB) Init() error {
 			Username:      username,
 			Password:      password,
 		}
-		s.clientOptions.SetAuth(credential)
+		s.options.SetAuth(credential)
 
 		// Check if TLS is enabled (via mongodb+srv:// or tls/ssl query params) and warn if not
 		parsedDSN, err := url.Parse(s.Dsn)
@@ -191,50 +166,73 @@ func (s *MongoDB) Init() error {
 			AuthSource:    "$external",
 			AuthMechanism: "MONGODB-X509",
 		}
-		s.clientOptions.SetAuth(credential)
+		s.options.SetAuth(credential)
 	default:
 		return fmt.Errorf("unsupported authentication type %q", s.AuthenticationType)
 	}
 
 	if s.ServerSelectTimeout != 0 {
-		s.clientOptions.SetServerSelectionTimeout(time.Duration(s.ServerSelectTimeout))
+		s.options.SetServerSelectionTimeout(time.Duration(s.ServerSelectTimeout))
 	}
 
-	s.clientOptions.ApplyURI(s.Dsn)
+	s.options.ApplyURI(s.Dsn)
 	return nil
 }
 
-func (s *MongoDB) createTimeSeriesCollection(databaseCollection string) error {
-	_, collectionExists := s.collections[databaseCollection]
-	if !collectionExists {
-		ctx := context.Background()
-		tso := options.TimeSeries()
-		tso.SetTimeField("timestamp")
-		tso.SetMetaField("tags")
-		tso.SetGranularity(s.MetricGranularity)
-		cco := options.CreateCollection()
-		if s.TTL != 0 {
-			cco.SetExpireAfterSeconds(int64(time.Duration(s.TTL).Seconds()))
+func (s *MongoDB) Connect() error {
+	// Connect to the database
+	ctx := context.Background()
+	client, err := mongo.Connect(ctx, s.options)
+	if err != nil {
+		return fmt.Errorf("connecting to server failed: %w", err)
+	}
+	s.client = client
+
+	// Cache the existing collections to prevent recreating those during write
+	collections, err := s.client.Database(s.MetricDatabase).ListCollections(ctx, bson.M{})
+	if err != nil {
+		return fmt.Errorf("listing collections failed: %w", err)
+	}
+
+	s.collections = make(map[string]bool, collections.RemainingBatchLength())
+	for collections.Next(ctx) {
+		var collection bson.M
+		if err = collections.Decode(&collection); err != nil {
+			return fmt.Errorf("decoding collections failed: %w", err)
 		}
-		cco.SetTimeSeriesOptions(tso)
-		err := s.client.Database(s.MetricDatabase).CreateCollection(ctx, databaseCollection, cco)
-		if err != nil {
-			return fmt.Errorf("unable to create time series collection: %w", err)
+
+		raw, found := collection["name"]
+		if !found {
+			return fmt.Errorf("name does not exist in collection %+v", collection)
+		}
+		name, ok := raw.(string)
+		if !ok {
+			return fmt.Errorf("non-string name %v (%T) in collection", raw, raw)
 		}
-		s.collections[databaseCollection] = bson.M{}
+		s.collections[name] = true
 	}
+
 	return nil
 }
 
-func (s *MongoDB) Connect() error {
+func (s *MongoDB) Write(metrics []telegraf.Metric) error {
 	ctx := context.Background()
-	client, err := mongo.Connect(ctx, s.clientOptions)
-	if err != nil {
-		return fmt.Errorf("unable to connect: %w", err)
-	}
-	s.client = client
-	if err = s.getCollections(ctx); err != nil {
-		return fmt.Errorf("unable to get collections from specified metric database: %w", err)
+
+	// Write one metric at a time
+	for _, metric := range metrics {
+		name := metric.Name()
+		// Create a new collection if it doesn't exist
+		if !s.collections[name] {
+			if err := s.createCollection(ctx, name); err != nil {
+				return fmt.Errorf("creating time series collection %q failed: %w", name, err)
+			}
+		}
+		doc := marshal(metric)
+
+		collection := s.client.Database(s.MetricDatabase).Collection(metric.Name())
+		if _, err := collection.InsertOne(ctx, &doc); err != nil {
+			return fmt.Errorf("getting collection %q failed: %w", metric.Name(), err)
+		}
 	}
 	return nil
 }
@@ -244,37 +242,45 @@ func (s *MongoDB) Close() error {
 	return s.client.Disconnect(ctx)
 }
 
-// all metric/measurement fields are parent level of document
-// metadata field is named "tags"
-// mongodb stores timestamp as UTC. conversion should be performed during reads in app or in aggregation pipeline
-func marshalMetric(metric telegraf.Metric) bson.D {
-	var bdoc bson.D
-	for k, v := range metric.Fields() {
-		bdoc = append(bdoc, primitive.E{Key: k, Value: v})
+func (s *MongoDB) createCollection(ctx context.Context, name string) error {
+	// Setup a new timeseries collection for the given metric name
+	series := options.TimeSeries()
+	series.SetTimeField("timestamp")
+	series.SetMetaField("tags")
+	series.SetGranularity(s.MetricGranularity)
+
+	collection := options.CreateCollection()
+	if s.TTL != 0 {
+		collection.SetExpireAfterSeconds(int64(time.Duration(s.TTL).Seconds()))
+	}
+	collection.SetTimeSeriesOptions(series)
+
+	// Create the new collection
+	if err := s.client.Database(s.MetricDatabase).CreateCollection(ctx, name, collection); err != nil {
+		return err
+	}
+	s.collections[name] = true
+
+	return nil
+}
+
+// Convert a metric into a MongoDB document with all fields being parent level
+// of document and the metadata field is named "tags. Mongodb stores timestamp
+// as UTC so conversion should be performed on the query or aggregation side.
+func marshal(metric telegraf.Metric) bson.D {
+	var doc bson.D
+	for _, f := range metric.FieldList() {
+		doc = append(doc, primitive.E{Key: f.Key, Value: f.Value})
 	}
 	var tags bson.D
-	for k, v := range metric.Tags() {
-		tags = append(tags, primitive.E{Key: k, Value: v})
+	for _, t := range metric.TagList() {
+		tags = append(tags, primitive.E{Key: t.Key, Value: t.Value})
 	}
-	bdoc = append(bdoc,
+	doc = append(doc,
 		primitive.E{Key: "tags", Value: tags},
 		primitive.E{Key: "timestamp", Value: metric.Time()},
 	)
-	return bdoc
-}
-
-func (s *MongoDB) Write(metrics []telegraf.Metric) error {
-	ctx := context.Background()
-	for _, metric := range metrics {
-		if err := s.createTimeSeriesCollection(metric.Name()); err != nil {
-			return err
-		}
-		bdoc := marshalMetric(metric)
-		if err := s.insertDocument(ctx, metric.Name(), bdoc); err != nil {
-			return err
-		}
-	}
-	return nil
+	return doc
 }
 
 func init() {

plugins/outputs/mongodb/mongodb_test.go

@@ -318,23 +318,18 @@ func TestConnectAndWriteIntegrationX509Auth(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			// validate config
-			err := tt.plugin.Init()
-			require.NoError(t, err)
+			require.NoError(t, tt.plugin.Init())
 
-			if err == nil {
-				// connect
-				err = tt.plugin.Connect()
-				tt.connErrFunc(t, err)
+			// connect
+			err = tt.plugin.Connect()
+			tt.connErrFunc(t, err)
 
-				if err == nil {
-					// insert mock metrics
-					err = tt.plugin.Write(testutil.MockMetrics())
-					require.NoError(t, err)
+			if err == nil {
+				// insert mock metrics
+				require.NoError(t, tt.plugin.Write(testutil.MockMetrics()))
 
-					// cleanup
-					err = tt.plugin.Close()
-					require.NoError(t, err)
-				}
+				// cleanup
+				require.NoError(t, tt.plugin.Close())
 			}
 		})
 	}