Integration of newrelic for mongodb collections

mongodb.go

package mongodb

import (
	"context"
	newrelic "github.com/newrelic/go-agent"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// Collection is a wrapper of the mongo.Collection type including a NewRelic Application dependency
type Collection struct {
	*mongo.Collection
	NrApp newrelic.Application
}

const (
	operationFind                   = "Find"
	operationBulkWrite              = "BulkWrite"
	operationInsertOne              = "InsertOne"
	operationInsertMany             = "InsertMany"
	operationDeleteOne              = "DeleteOne"
	operationDeleteMany             = "DeleteMany"
	operationUpdateOne              = "UpdateOne"
	operationUpdateMany             = "UpdateMany"
	operationReplaceOne             = "ReplaceOne"
	operationAggregate              = "Aggregate"
	operationCountDocuments         = "CountDocuments"
	operationEstimatedDocumentCount = "EstimatedDocumentCount"
	operationDistinct               = "Distinct"
	operationFindOne                = "FindOne"
	operationFindOneAndDelete       = "FindOneAndDelete"
	operationFindOneAndReplace      = "FindOneAndReplace"
	operationFindOneAndUpdate       = "FindOneAndUpdate"
	operationDrop                   = "Drop"
)

func buildDatastoreSegment(coll *Collection, operation string) newrelic.DatastoreSegment {
	return newrelic.DatastoreSegment{
		Product:    newrelic.DatastoreMongoDB,
		Collection: coll.Name(),
		Operation:  operation,
		DatabaseName: coll.Database().Name(),
	}
}

// Clone creates a copy of this collection with updated options, if any are given.
func (coll *Collection) Clone(opts ...*options.CollectionOptions) (*Collection, error) {
	cl, err := coll.Collection.Clone(opts...)
	return &Collection{
		Collection: cl,
		NrApp:      coll.NrApp,
	}, err

}

// Find finds the documents matching a model.
func (coll *Collection) Find(ctx context.Context, filter interface{},
	opts ...*options.FindOptions) (*mongo.Cursor, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationFind, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationFind)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.Find(ctx, filter, opts...)
}

// BulkWrite performs a bulk write operation.
//
// See https://docs.mongodb.com/manual/core/bulk-write-operations/.
func (coll *Collection) BulkWrite(ctx context.Context, models []mongo.WriteModel,
	opts ...*options.BulkWriteOptions) (*mongo.BulkWriteResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationBulkWrite, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationBulkWrite)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.BulkWrite(ctx, models, opts...)
}

// InsertOne inserts a single document into the collection.
func (coll *Collection) InsertOne(ctx context.Context, document interface{},
	opts ...*options.InsertOneOptions) (*mongo.InsertOneResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationInsertOne, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationInsertOne)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.InsertOne(ctx, document, opts...)
}

// InsertMany inserts the provided documents.
func (coll *Collection) InsertMany(ctx context.Context, documents []interface{},
	opts ...*options.InsertManyOptions) (*mongo.InsertManyResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationInsertMany, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationInsertMany)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.InsertMany(ctx, documents, opts...)
}

// DeleteOne deletes a single document from the collection.
func (coll *Collection) DeleteOne(ctx context.Context, filter interface{},
	opts ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationDeleteOne, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationDeleteOne)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.DeleteOne(ctx, filter, opts...)
}

// DeleteMany deletes multiple documents from the collection.
func (coll *Collection) DeleteMany(ctx context.Context, filter interface{},
	opts ...*options.DeleteOptions) (*mongo.DeleteResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationDeleteMany, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationDeleteMany)
	defer segment.End() // nolint:errcheck
	return coll.Collection.DeleteMany(ctx, filter, opts...)
}

// UpdateOne updates a single document in the collection.
func (coll *Collection) UpdateOne(ctx context.Context, filter interface{}, update interface{},
	opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationUpdateOne, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationUpdateOne)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.UpdateOne(ctx, filter, update, opts...)
}

// UpdateMany updates multiple documents in the collection.
func (coll *Collection) UpdateMany(ctx context.Context, filter interface{}, update interface{},
	opts ...*options.UpdateOptions) (*mongo.UpdateResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationUpdateMany, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationUpdateMany)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.UpdateMany(ctx, filter, update, opts...)
}

// ReplaceOne replaces a single document in the collection.
func (coll *Collection) ReplaceOne(ctx context.Context, filter interface{}, replacement interface{},
	opts ...*options.ReplaceOptions) (*mongo.UpdateResult, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationReplaceOne, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationReplaceOne)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.ReplaceOne(ctx, filter, replacement, opts...)
}

// Aggregate runs an aggregation framework pipeline.
//
// See https://docs.mongodb.com/manual/aggregation/.
func (coll *Collection) Aggregate(ctx context.Context, pipeline interface{},
	opts ...*options.AggregateOptions) (*mongo.Cursor, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationAggregate, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationAggregate)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.Aggregate(ctx, pipeline, opts...)
}

// CountDocuments gets the number of documents matching the filter.
func (coll *Collection) CountDocuments(ctx context.Context, filter interface{},
	opts ...*options.CountOptions) (int64, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationCountDocuments, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationCountDocuments)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.CountDocuments(ctx, filter, opts...)
}

// EstimatedDocumentCount gets an estimate of the count of documents in a collection using collection metadata.
func (coll *Collection) EstimatedDocumentCount(ctx context.Context,
	opts ...*options.EstimatedDocumentCountOptions) (int64, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationEstimatedDocumentCount, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationEstimatedDocumentCount)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.EstimatedDocumentCount(ctx, opts...)
}

// Distinct finds the distinct values for a specified field across a single
// collection.
func (coll *Collection) Distinct(ctx context.Context, fieldName string, filter interface{},
	opts ...*options.DistinctOptions) ([]interface{}, error) {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationDistinct, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationDistinct)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.Distinct(ctx, fieldName, filter, opts...)
}

// FindOne returns up to one document that matches the model.
func (coll *Collection) FindOne(ctx context.Context, filter interface{},
	opts ...*options.FindOneOptions) *mongo.SingleResult {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationFindOne, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationFindOne)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.FindOne(ctx, filter, opts...)
}

// FindOneAndDelete find a single document and deletes it, returning the
// original in result.
func (coll *Collection) FindOneAndDelete(ctx context.Context, filter interface{},
	opts ...*options.FindOneAndDeleteOptions) *mongo.SingleResult {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationFindOneAndDelete, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationFindOneAndDelete)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.FindOneAndDelete(ctx, filter, opts...)
}

// FindOneAndReplace finds a single document and replaces it, returning either
// the original or the replaced document.
func (coll *Collection) FindOneAndReplace(ctx context.Context, filter interface{}, replacement interface{},
	opts ...*options.FindOneAndReplaceOptions) *mongo.SingleResult {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationFindOneAndReplace, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationFindOneAndReplace)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.FindOneAndReplace(ctx, filter, replacement, opts...)
}

// FindOneAndUpdate finds a single document and updates it, returning either
// the original or the updated.
func (coll *Collection) FindOneAndUpdate(ctx context.Context, filter interface{}, update interface{},
	opts ...*options.FindOneAndUpdateOptions) *mongo.SingleResult {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationFindOneAndUpdate, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationFindOneAndUpdate)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.FindOneAndUpdate(ctx, filter, update, opts...)
}

// Drop drops this collection from database.
func (coll *Collection) Drop(ctx context.Context) error {
	if newrelic.FromContext(ctx) == nil {
		txn := coll.NrApp.StartTransaction(operationDrop, nil, nil)
		defer txn.End() // nolint:errcheck
		ctx = newrelic.NewContext(ctx, txn)
	}

	segment := buildDatastoreSegment(coll, operationDrop)
	defer segment.End() // nolint:errcheck

	segment.StartTime = newrelic.StartSegmentNow(newrelic.FromContext(ctx))
	return coll.Collection.Drop(ctx)
}