Merge pull request #24 from GabrielCalin/master

Add support for other Mongo operations

README.md

@@ -8,6 +8,10 @@ K6 extension to perform tests on mongo.
 - Supports inserting document batch.
 - Supports find a document based on filter.
 - Supports find all documents of a collection.
+- Supports upserting a document based on filter.
+- Supports bulk upserting documents based on filters.
+- Supports aggregation pipelines.
+- Supports finding distinct values for a field in a collection based on a filter.
 - Supports delete first document based on filter.
 - Supports deleting all documents for a specific filter.
 - Supports dropping a collection.

go.mod

@@ -4,10 +4,11 @@ go 1.19
 
 require (
 	go.k6.io/k6 v0.45.1
-	go.mongodb.org/mongo-driver v1.9.1
+	go.mongodb.org/mongo-driver v1.13.1
 )
 
 require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/dlclark/regexp2 v1.9.0 // indirect
 	github.com/dop251/goja v0.0.0-20230531210528-d7324b2d74f7 // indirect
 	github.com/fatih/color v1.15.0 // indirect
@@ -21,6 +22,7 @@ require (
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.18 // indirect
+	github.com/montanaflynn/stats v0.0.0-20171201202039-1bf9dbcd8cbe // indirect
 	github.com/mstoykov/atlas v0.0.0-20220811071828-388f114305dd // indirect
 	github.com/onsi/ginkgo v1.16.5 // indirect
 	github.com/onsi/gomega v1.27.10 // indirect
@@ -29,8 +31,8 @@ require (
 	github.com/sirupsen/logrus v1.9.0 // indirect
 	github.com/spf13/afero v1.8.2 // indirect
 	github.com/xdg-go/pbkdf2 v1.0.0 // indirect
-	github.com/xdg-go/scram v1.1.1 // indirect
-	github.com/xdg-go/stringprep v1.0.3 // indirect
+	github.com/xdg-go/scram v1.1.2 // indirect
+	github.com/xdg-go/stringprep v1.0.4 // indirect
 	github.com/youmark/pkcs8 v0.0.0-20201027041543-1326539a0a0a // indirect
 	golang.org/x/crypto v0.17.0 // indirect
 	golang.org/x/sync v0.2.0 // indirect

mongo.go

@@ -2,11 +2,11 @@ package xk6_mongo
 
 import (
 	"context"
-	"log"
-
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"log"
+	"log/slog"
 
 	k6modules "go.k6.io/k6/js/modules"
 )
@@ -25,27 +25,36 @@ type Client struct {
 	client *mongo.Client
 }
 
+type UpsertOneModel struct {
+	Query  interface{} `json:"query"`
+	Update interface{} `json:"update"`
+}
+
 // NewClient represents the Client constructor (i.e. `new mongo.Client()`) and
 // returns a new Mongo client object.
 // connURI -> mongodb://username:password@address:port/db?connect=direct
 func (*Mongo) NewClient(connURI string) interface{} {
+	log.Print("start creating new client")
 
 	clientOptions := options.Client().ApplyURI(connURI)
 	client, err := mongo.Connect(context.TODO(), clientOptions)
 	if err != nil {
+		slog.Error(err.Error())
 		return err
 	}
 
+	log.Print("created new client")
 	return &Client{client: client}
 }
 
 const filter_is string = "filter is "
 
-func (c *Client) Insert(database string, collection string, doc map[string]string) error {
+func (c *Client) Insert(database string, collection string, doc interface{}) error {
 	db := c.client.Database(database)
 	col := db.Collection(collection)
 	_, err := col.InsertOne(context.TODO(), doc)
 	if err != nil {
+		slog.Error(err.Error())
 		return err
 	}
 	return nil
@@ -57,22 +66,68 @@ func (c *Client) InsertMany(database string, collection string, docs []any) erro
 	col := db.Collection(collection)
 	_, err := col.InsertMany(context.TODO(), docs)
 	if err != nil {
+		slog.Error(err.Error())
 		return err
 	}
 	return nil
 }
 
-func (c *Client) Find(database string, collection string, filter interface{}) []bson.M {
+func (c *Client) Upsert(database string, collection string, filter interface{}, upsert interface{}) error {
+	db := c.client.Database(database)
+	col := db.Collection(collection)
+	opts := options.Update().SetUpsert(true)
+	_, err := col.UpdateOne(context.TODO(), filter, upsert, opts)
+	if err != nil {
+		slog.Error(err.Error())
+	}
+	return nil
+}
+
+func (c *Client) BulkUpsert(database string, collection string, upserts []UpsertOneModel) error {
 	db := c.client.Database(database)
 	col := db.Collection(collection)
+	var models []mongo.WriteModel
+	for _, upsert := range upserts {
+		model := mongo.NewUpdateOneModel()
+		model.SetFilter(upsert.Query)
+		model.SetUpdate(upsert.Update)
+		model.SetUpsert(true)
+		models = append(models, model)
+	}
+	_, err := col.BulkWrite(context.TODO(), models)
+	if err != nil {
+		slog.Error(err.Error())
+	}
+	return nil
+}
+
+func (c *Client) Find(database string, collection string, filter interface{}, sort interface{}, limit int64) []bson.M {
+	db := c.client.Database(database)
+	col := db.Collection(collection)
+	opts := options.Find().SetSort(sort).SetLimit(limit)
 	log.Print(filter_is, filter)
-	cur, err := col.Find(context.TODO(), filter)
+	cur, err := col.Find(context.TODO(), filter, opts)
 	if err != nil {
-		log.Fatal(err)
+		slog.Error(err.Error())
+	}
+	var results []bson.M
+	if err = cur.All(context.TODO(), &results); err != nil {
+		slog.Error(err.Error())
+	}
+	return results
+}
+
+func (c *Client) Aggregate(database string, collection string, pipeline interface{}) []bson.M {
+	db := c.client.Database(database)
+	col := db.Collection(collection)
+	log.Print(filter_is, pipeline)
+	cur, err := col.Aggregate(context.TODO(), pipeline)
+	if err != nil {
+		slog.Error(err.Error())
 	}
 	var results []bson.M
 	if err = cur.All(context.TODO(), &results); err != nil {
-		panic(err)
+		slog.Error(err.Error())
 	}
 	return results
 }
@@ -89,7 +144,7 @@ func (c *Client) FindOne(database string, collection string, filter map[string]s
 		return nil
 	}
 	if err != nil {
-		log.Fatal(err)
+		slog.Error(err.Error())
 	}
 	log.Printf("found document %v", result)
 	return nil
@@ -121,23 +176,22 @@ func (c *Client) FindAll(database string, collection string) []bson.M {
 	col := db.Collection(collection)
 	cur, err := col.Find(context.TODO(), bson.D{{}})
 	if err != nil {
-		log.Fatal(err)
+		slog.Error(err.Error())
 	}
 	var results []bson.M
 	if err = cur.All(context.TODO(), &results); err != nil {
-		panic(err)
+		slog.Error(err.Error())
 	}
 	return results
 }
 
 func (c *Client) DeleteOne(database string, collection string, filter map[string]string) error {
 	db := c.client.Database(database)
 	col := db.Collection(collection)
-	opts := options.Delete().SetHint(bson.D{{"_id", 1}})
 	log.Print(filter_is, filter)
-	result, err := col.DeleteOne(context.TODO(), filter, opts)
+	result, err := col.DeleteOne(context.TODO(), filter)
 	if err != nil {
-		log.Fatal(err)
+		slog.Error(err.Error())
 	}
 	log.Printf("Deleted documents %v", result)
 	return nil
@@ -146,23 +200,41 @@ func (c *Client) DeleteOne(database string, collection string, filter map[string
 func (c *Client) DeleteMany(database string, collection string, filter map[string]string) error {
 	db := c.client.Database(database)
 	col := db.Collection(collection)
-	opts := options.Delete().SetHint(bson.D{{"_id", 1}})
 	log.Print(filter_is, filter)
-	result, err := col.DeleteMany(context.TODO(), filter, opts)
+	result, err := col.DeleteMany(context.TODO(), filter)
 	if err != nil {
-		log.Fatal(err)
+		slog.Error(err.Error())
 	}
 	log.Printf("Deleted documents %v", result)
 	return nil
 }
 
+func (c *Client) Distinct(database string, collection string, field string, filter interface{}) []interface{} {
+	db := c.client.Database(database)
+	col := db.Collection(collection)
+	results, err := col.Distinct(context.TODO(), field, filter)
+	if err != nil {
+		slog.Error(err.Error())
+	}
+
+	return results
+}
+
 func (c *Client) DropCollection(database string, collection string) error {
 	log.Printf("Delete collection if present")
 	db := c.client.Database(database)
 	col := db.Collection(collection)
 	err := col.Drop(context.TODO())
 	if err != nil {
-		log.Fatal(err)
+		slog.Error(err.Error())
 	}
 	return nil
 }
+
+func (c *Client) Disconnect() {
+	log.Printf("Disconnecting from Mongo database")
+	err := c.client.Disconnect(context.TODO())
+	if err != nil {
+		slog.Error(err.Error())
+	}
+}