mongo connection

Author: OlgaMinyaeva

.travis.yml

@@ -10,6 +10,7 @@ after_success:
   - bash <(curl -s https://codecov.io/bash)
 services:
   - mysql
+  - mongodb
 
 env:
-  - TEST_MYSQL_CONNECTION="root:@(127.0.0.1:3306)/%dbname%?tx_isolation=SERIALIZABLE&parseTime=true"
+  - TEST_MYSQL_CONNECTION="root:@(127.0.0.1:3306)/%dbname%?tx_isolation=SERIALIZABLE&parseTime=true";TEST_MONGO_CONNECTION="mongodb://localhost:27017"

README.md

@@ -60,17 +60,18 @@ Configurations are stored in json files. Example:
 | logger.level  | `debug`, `info`, `warning`, `error` | Logging level.         |
 
 ### DB - Configuration of storage.
-| Option            | Possible values  | Description                                |
-| ----------------- | ---------------  | ------------------------------------------ |
-| db.implementation | `mysql`, `local` | Select your favorite db, or local storage. |
-| db.connection     | see next table   | DSN for your db.                           |
+| Option            | Possible values          | Description                                |
+| ----------------- | ------------------------ | ------------------------------------------ |
+| db.implementation | `mysql`, `local`, `mongo`| Select your favorite db, or local storage. |
+| db.connection     | see next table           | DSN for your db.                           |
 
 >Note: Note: Local (in-memory) storage does not support common session storage between grid-instances.
 
 | DB implementation | DSN format                                                                                                                                                              |
 | ----------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | mysql             | [spec](https://github.com/go-sql-driver/mysql#dsn-data-source-name), example `db_user:db_pass@(db_host:3306)/db_name?parseTime=true` (parseTime=true - required option) |
 | local             | omit this property, because every instance have its own in-memory storage                                                                                               |
+| mongo             | NOTE! Mongo db temporary supports only persistent node strategy, example `mongodb://localhost:27017`
 
 ### Statsd - Configuration of metrics(optional).
 | Option          | Possible values | Description            |

config/config.go

@@ -53,6 +53,7 @@ type Logger struct {
 type DB struct {
 	Implementation string `json:"implementation"`
 	Connection     string `json:"connection"`
+	DbName         string `json:"db_name"`
 }
 
 // Statsd - Settings of metrics sender.

handlers/createSession.go

@@ -74,7 +74,7 @@ func (h *CreateSession) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
 func (h *CreateSession) tryCreateSession(r *http.Request, capabilities *capabilities.Capabilities) (*proxy.ResponseWriter, error) {
 	select {
 	case <-r.Context().Done():
-		err := errors.New("Request cancelled by client, " + r.Context().Err().Error())
+		err := errors.New("Request canceled by client, " + r.Context().Err().Error())
 		return nil, err
 	default:
 	}

invoke.go

@@ -3,6 +3,8 @@ package main
 import (
 	_ "github.com/go-sql-driver/mysql"
 
+	"github.com/qa-dev/jsonwire-grid/storage/mongo"
+
 	"github.com/qa-dev/jsonwire-grid/config"
 
 	"errors"
@@ -32,6 +34,13 @@ func invokeStorageFactory(config config.Config) (factory StorageFactoryInterface
 		factory = new(mysql.Factory)
 	case "local":
 		factory = new(local.Factory)
+	case "mongo":
+		for _, s := range config.Grid.StrategyList {
+			if s.Type != string(pool.NodeTypePersistent) {
+				err = errors.New("Invalid config: Mongo db supports only persistent node strategy ")
+			}
+		}
+		factory = new(mongo.Factory)
 	default:
 		err = errors.New("Invalid config, unknown param [db.implementation=" + config.DB.Implementation + "]")
 	}

pool/node.go

@@ -22,14 +22,14 @@ type Node struct {
 	// The value may depend on the strategy:
 	// - for constant nodes ip: port
 	// - for temporary pod.name
-	Key              string                      `json:"key"`
-	Type             NodeType                    `json:"type"`
-	Address          string                      `json:"address"`
-	Status           NodeStatus                  `json:"status"`
-	SessionID        string                      `json:"session_id"`
-	Updated          int64                       `json:"updated"`
-	Registered       int64                       `json:"registered"`
-	CapabilitiesList []capabilities.Capabilities `json:"capabilities_list"`
+	Key              string                      `json:"key" bson:"key"`
+	Type             NodeType                    `json:"type" bson:"type"`
+	Address          string                      `json:"address" bson:"address"`
+	Status           NodeStatus                  `json:"status" bson:"status"`
+	SessionID        string                      `json:"session_id" bson:"session_id"`
+	Updated          int64                       `json:"updated" bson:"updated"`
+	Registered       int64                       `json:"registered" bson:"registered"`
+	CapabilitiesList []capabilities.Capabilities `json:"capabilities_list" bson:"capabilities_list"`
 }
 
 func (n *Node) String() string {

pool/strategy/kubernetes/strategy.go

@@ -48,7 +48,6 @@ func (s *Strategy) Reserve(desiredCaps capabilities.Capabilities) (pool.Node, er
 	}
 	node.Address = nodeAddress
 	return *node, nil
-
 }
 
 func (s *Strategy) CleanUp(node pool.Node) error {

pool/strategy/persistent/strategy.go

@@ -112,7 +112,6 @@ func (s *Strategy) registerCapabilities(nodeList []pool.Node) {
 	for _, node := range nodeList {
 		for _, availableCaps := range node.CapabilitiesList {
 			s.capsComparator.Register(availableCaps)
-
 		}
 	}
 }

storage/mongo/factory.go

@@ -0,0 +1,47 @@
+package mongo
+
+import (
+	"context"
+	"errors"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+
+	"github.com/qa-dev/jsonwire-grid/config"
+	"github.com/qa-dev/jsonwire-grid/pool"
+)
+
+type Factory struct {
+}
+
+func (f *Factory) Create(cfg config.Config) (pool.StorageInterface, error) {
+	ctx := context.Background()
+	client, err := mongo.Connect(ctx, options.Client().ApplyURI(cfg.DB.Connection))
+	if err != nil {
+		panic("Database connection error: " + err.Error())
+	}
+
+	err = client.Ping(ctx, nil)
+	if err != nil {
+		err = errors.New("Database connection not established: " + err.Error())
+		return nil, err
+	}
+
+	db := client.Database(cfg.DB.DbName)
+	s := NewMongoStorage(db)
+	mod := mongo.IndexModel{
+		Keys: bson.M{
+			"key":     1,
+			"address": 1,
+		},
+		Options: options.Index().SetUnique(true).SetName("key_address_unique"),
+	}
+
+	_, err = s.collection.Indexes().CreateOne(ctx, mod)
+	if err != nil {
+		err = errors.New("Create index error: " + err.Error())
+		return nil, err
+	}
+	return s, nil
+}

storage/mongo/mongo.go

@@ -0,0 +1,194 @@
+package mongo
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+
+	"github.com/qa-dev/jsonwire-grid/pool"
+	"github.com/qa-dev/jsonwire-grid/storage"
+)
+
+type Storage struct {
+	collection *mongo.Collection
+	ctx        context.Context
+}
+
+func NewMongoStorage(db *mongo.Database) *Storage {
+	return &Storage{
+		collection: db.Collection("grid"),
+		ctx:        context.Background(),
+	}
+}
+
+func (m *Storage) Add(node pool.Node, limit int) error {
+	opts := options.Update().SetUpsert(true)
+	filter := bson.M{"key": node.Key}
+	update := bson.D{{Key: "$set", Value: node}}
+
+	if limit > 0 {
+		return errors.New("[Mongo/Add] limit strategy temporary not supported ")
+	}
+
+	result, err := m.collection.UpdateOne(m.ctx, filter, update, opts)
+	if err != nil {
+		return errors.New("[Mongo/Add] update node, " + err.Error())
+	}
+	if result.MatchedCount == 0 && result.UpsertedCount == 0 {
+		return errors.New("[Mongo/Add] No rows was affected ")
+	}
+
+	return nil
+}
+
+func (m *Storage) ReserveAvailable(nodeList []pool.Node) (pool.Node, error) {
+	nodeKeyList := make([]string, 0, len(nodeList))
+	node := pool.Node{}
+	for _, node := range nodeList {
+		nodeKeyList = append(nodeKeyList, node.Key)
+	}
+
+	filter := bson.M{"key": bson.M{"$in": nodeKeyList}, "status": pool.NodeStatusAvailable}
+
+	update := bson.M{"$set": bson.M{
+		"updated": time.Now().Unix(),
+		"status":  string(pool.NodeStatusReserved),
+	}}
+
+	opts := options.
+		FindOneAndUpdate().
+		SetReturnDocument(options.After).
+		SetSort(bson.M{"updated": 1})
+
+	err := m.collection.FindOneAndUpdate(m.ctx, filter, update, opts).Decode(&node)
+	if err != nil {
+		err = errors.New("[Mongo/ReserveAvailable] find and update node, " + err.Error())
+		return node, err
+	}
+	return node, nil
+}
+
+func (m *Storage) SetBusy(node pool.Node, sessionID string) error {
+	filter := bson.M{"key": node.Key}
+	update := bson.M{"$set": bson.M{
+		"session_id": sessionID,
+		"updated":    time.Now().Unix(),
+		"status":     string(pool.NodeStatusBusy),
+	}}
+	result, err := m.collection.UpdateOne(m.ctx, filter, update)
+	if err != nil {
+		err = errors.New("[Mongo/SetBusy] update node in collection, " + err.Error())
+		return err
+	}
+	if result.ModifiedCount == 0 {
+		return storage.ErrNotFound
+	}
+	return nil
+}
+
+func (m *Storage) SetAvailable(node pool.Node) error {
+	filter := bson.M{"key": node.Key}
+	update := bson.M{"$set": bson.M{
+		"updated": time.Now().Unix(),
+		"status":  string(pool.NodeStatusAvailable),
+	}}
+	result, err := m.collection.UpdateOne(m.ctx, filter, update)
+	if err != nil {
+		err = errors.New("[Mongo/SetAvailable] update node in collection, " + err.Error())
+		return err
+	}
+	if result.ModifiedCount == 0 {
+		return storage.ErrNotFound
+	}
+	return nil
+}
+
+func (m *Storage) GetCountWithStatus(status *pool.NodeStatus) (int, error) {
+	var count int
+	var err error
+	filter := bson.M{}
+	if status != nil {
+		filter = bson.M{"status": status}
+	}
+	cnt, err := m.collection.CountDocuments(m.ctx, filter)
+	count = int(cnt)
+	if err != nil {
+		return 0, errors.New("[Mongo/GetCountWithStatus] count nodes in collection, " + err.Error())
+	}
+	return count, nil
+}
+
+func (m *Storage) GetBySession(session string) (pool.Node, error) {
+	node, err := m.getByField("session_id", session)
+	if err != nil {
+		return pool.Node{}, errors.New("[Mongo/GetBySession] find node in collection, " + err.Error())
+	}
+	return node, nil
+}
+
+func (m *Storage) GetByAddress(address string) (pool.Node, error) {
+	node, err := m.getByField("address", address)
+	if err != nil {
+		return pool.Node{}, errors.New("[Mongo/GetByAddress] find node in collection, " + err.Error())
+	}
+	return node, nil
+}
+
+func (m *Storage) GetAll() ([]pool.Node, error) {
+	nodeList := make([]pool.Node, 0)
+	resultCursor, err := m.collection.Find(m.ctx, bson.M{})
+	if err != nil {
+		return nil, errors.New("[Mongo/GetAll] find node in collection, " + err.Error())
+	}
+	defer resultCursor.Close(m.ctx)
+	for resultCursor.Next(m.ctx) {
+		var result pool.Node
+		err := resultCursor.Decode(&result)
+		if err != nil {
+			return nil, errors.New("[Mongo/GetAll] decode node data" + err.Error())
+		}
+		nodeList = append(nodeList, result)
+	}
+	if err := resultCursor.Err(); err != nil {
+		return nil, errors.New("[Mongo/GetAll] iterate result" + err.Error())
+	}
+	return nodeList, nil
+}
+
+func (m *Storage) Remove(node pool.Node) error {
+	rowsAffected, err := m.collection.DeleteOne(m.ctx, bson.M{"key": node.Key})
+	if err != nil {
+		return errors.New("[Mongo/Remove] delete from node collection, " + err.Error())
+	}
+	if rowsAffected.DeletedCount == 0 {
+		return errors.New("[Mongo/Remove] delete from node collection: affected 0 rows")
+	}
+	return nil
+}
+
+func (m *Storage) UpdateAddress(node pool.Node, newAddress string) error {
+	filter := bson.M{"key": node.Key}
+	update := bson.M{"$set": bson.M{"address": newAddress}}
+	result, err := m.collection.UpdateOne(m.ctx, filter, update)
+	if err != nil {
+		err = errors.New("[Mongo/UpdateAddress], " + err.Error())
+		return err
+	}
+	if result.ModifiedCount == 0 {
+		return storage.ErrNotFound
+	}
+	return nil
+}
+
+func (m *Storage) getByField(key, value string) (pool.Node, error) {
+	node := pool.Node{}
+	filter := bson.M{
+		key: bson.M{"$eq": value},
+	}
+	err := m.collection.FindOne(m.ctx, filter).Decode(&node)
+	return node, err
+}