/
mongo.go
99 lines (84 loc) · 2.4 KB
/
mongo.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
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"
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
}
// 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{} {
clientOptions := options.Client().ApplyURI(connURI)
client, err := mongo.Connect(context.TODO(), clientOptions)
if err != nil {
return err
}
return &Client{client: client}
}
func (c *Client) DropDatabase(database string) error {
err := c.client.Database(database).Drop(context.TODO())
if err != nil {
return err
}
return nil
}
func (c *Client) InsertOne(database string, collection string, doc map[string]string) error {
db := c.client.Database(database)
col := db.Collection(collection)
_, err := col.InsertOne(context.TODO(), doc)
if err != nil {
return err
}
return nil
}
func (c *Client) DeleteOne(database string, collection string, filter interface{}) error {
db := c.client.Database(database)
col := db.Collection(collection)
_, err := col.DeleteOne(context.TODO(), filter)
if err != nil {
return err
}
return nil
}
func (c *Client) Find(database string, collection string, filter interface{}) []bson.M {
db := c.client.Database(database)
col := db.Collection(collection)
log.Print("filter is ", filter)
cur, err := col.Find(context.TODO(), filter)
if err != nil {
log.Fatal(err)
// return nil
}
var results []bson.M
if err = cur.All(context.TODO(), &results); err != nil {
panic(err)
}
return results
}
func (c *Client) FindOne(database string, collection string, filter map[string]string) error {
db := c.client.Database(database)
col := db.Collection(collection)
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 != nil {
log.Fatal(err)
}
log.Printf("found document %v", result)
return nil
}