mongo.go

import (
        "context"
        "log"
        "time"

        "go.mongodb.org/mongo-driver/bson"
        "go.mongodb.org/mongo-driver/mongo"
        "go.mongodb.org/mongo-driver/mongo/options"

        k6modules "go.k6.io/k6/js/modules"
)

// Register the extension on module initialization, available to
// import from JS as "k6/x/mongo".
func init() {
        k6modules.Register("k6/x/mongo", new(Mongo))
}

// Mongo is the k6 extension for a Mongo client.
type Mongo struct{}

// Client is the Mongo client wrapper.
type Client struct {
        client *mongo.Client
}

// Helper: safely close Mongo connection
func (c *Client) safeClose() {
        ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
        defer cancel()
        if err := c.client.Disconnect(ctx); err != nil {
                log.Printf("Error closing Mongo client: %v", err)
        }
}

// NewClient creates a new Mongo client connection
// connURI -> mongodb://username:password@address:port/db?connect=direct
func (*Mongo) NewClient(connURI string) interface{} {
        clientOptions := options.Client().ApplyURI(connURI)
        client, err := mongo.Connect(context.TODO(), clientOptions)
        if err != nil {
                return err
        }
        return &Client{client: client}
}

const filter_is string = "filter is "

// Insert one document
func (c *Client) Insert(database string, collection string, doc interface{}) error {
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()
        _, err := col.InsertOne(context.TODO(), doc)
        return err
}

// Insert multiple documents
func (c *Client) InsertMany(database string, collection string, docs []any) error {
        log.Printf("Insert multiple documents")
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()
        _, err := col.InsertMany(context.TODO(), docs)
        return err
}

// Find many documents
func (c *Client) Find(database string, collection string, filter interface{}) []bson.M {
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()
        log.Print(filter_is, filter)

        cur, err := col.Find(context.TODO(), filter)
        if err != nil {
                log.Fatal(err)
        }
        var results []bson.M
        if err = cur.All(context.TODO(), &results); err != nil {
                panic(err)
        }
        return results
}

// Find one document
func (c *Client) FindOne(database string, collection string, filter interface{}) error {
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()

        var result bson.M
        opts := options.FindOne().SetSort(bson.D{{"_id", 1}})
        log.Print(filter_is, filter)
        err := col.FindOne(context.TODO(), filter, opts).Decode(&result)
        if err == mongo.ErrNoDocuments {
                log.Printf("No document was found for filter %v", filter)
                return nil
        }
        if err != nil {
                log.Fatal(err)
        }
        log.Printf("found document %v", result)
        return nil
}

// Update one document
func (c *Client) UpdateOne(database string, collection string, filter interface{}, data interface{}) error {
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()

        update := bson.D{{"$set", data}}
        result, err := col.UpdateOne(context.TODO(), filter, update)
        if err != nil {
                log.Fatal(err)
        }
        log.Printf("updated document %v", result)
        return nil
}

// Find all documents
func (c *Client) FindAll(database string, collection string) []bson.M {
        log.Printf("Find all documents")
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()

        cur, err := col.Find(context.TODO(), bson.D{{}})
        if err != nil {
                log.Fatal(err)
        }
        var results []bson.M
        if err = cur.All(context.TODO(), &results); err != nil {
                panic(err)
        }
        return results
}

// Delete one document
func (c *Client) DeleteOne(database string, collection string, filter interface{}) error {
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()

        opts := options.Delete().SetHint(bson.D{{"_id", 1}})
        log.Print(filter_is, filter)
        result, err := col.DeleteOne(context.TODO(), filter, opts)
        if err != nil {
                log.Fatal(err)
        }
        log.Printf("Deleted documents %v", result)
        return nil
}

// Delete many documents
func (c *Client) DeleteMany(database string, collection string, filter interface{}) error {
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()

        opts := options.Delete().SetHint(bson.D{{"_id", 1}})
        log.Print(filter_is, filter)
        result, err := col.DeleteMany(context.TODO(), filter, opts)
        if err != nil {
                log.Fatal(err)
        }
        log.Printf("Deleted documents %v", result)
        return nil
}

// Drop a collection
func (c *Client) DropCollection(database string, collection string) error {
        log.Printf("Delete collection if present")
        db := c.client.Database(database)
        col := db.Collection(collection)
        defer c.safeClose()

        err := col.Drop(context.TODO())
        if err != nil {
                log.Fatal(err)
        }
        return nil
}