1 files changed, 78 insertions, 9 deletions

diff --git a/storage.go b/storage.go
index aced06b..8505bfe 100644
--- a/storage.go
+++ b/storage.go
@@ -1,3 +1,20 @@
+/*
+ * This file is part of caronte (https://github.com/eciavatta/caronte).
+ * Copyright (c) 2020 Emiliano Ciavatta.
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, version 3.
+ *
+ * This program is distributed in the hope that it will be useful, but
+ * WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
+ * General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
 package main
 
 import (
@@ -8,15 +25,20 @@ import (
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
 	"go.mongodb.org/mongo-driver/mongo/options"
+	"time"
 )
 
 // Collections names
-const Connections = "connections"
-const ConnectionStreams = "connection_streams"
-const ImportingSessions = "importing_sessions"
-const Rules = "rules"
-const Settings = "settings"
-const Services = "services"
+const (
+	Connections       = "connections"
+	ConnectionStreams = "connection_streams"
+	ImportingSessions = "importing_sessions"
+	Rules             = "rules"
+	Searches          = "searches"
+	Settings          = "settings"
+	Services          = "services"
+	Statistics        = "statistics"
+)
 
 var ZeroRowID [12]byte
 
@@ -55,8 +77,10 @@ func NewMongoStorage(uri string, port int, database string) (*MongoStorage, erro
 		ConnectionStreams: db.Collection(ConnectionStreams),
 		ImportingSessions: db.Collection(ImportingSessions),
 		Rules:             db.Collection(Rules),
+		Searches:          db.Collection(Searches),
 		Settings:          db.Collection(Settings),
 		Services:          db.Collection(Services),
+		Statistics:        db.Collection(Statistics),
 	}
 
 	if _, err := collections[Services].Indexes().CreateOne(ctx, mongo.IndexModel{
@@ -66,9 +90,13 @@ func NewMongoStorage(uri string, port int, database string) (*MongoStorage, erro
 		return nil, err
 	}
 
-	if _, err := collections[ConnectionStreams].Indexes().CreateOne(ctx, mongo.IndexModel{
-		Keys:    bson.D{{"connection_id", -1}}, // descending
-		Options: options.Index(),
+	if _, err := collections[ConnectionStreams].Indexes().CreateMany(ctx, []mongo.IndexModel{
+		{
+			Keys: bson.D{{"connection_id", -1}}, // descending
+		},
+		{
+			Keys: bson.D{{"payload_string", "text"}},
+		},
 	}); err != nil {
 		return nil, err
 	}
@@ -150,6 +178,7 @@ type UpdateOperation interface {
 	Filter(filter OrderedDocument) UpdateOperation
 	Upsert(upsertResults *interface{}) UpdateOperation
 	One(update interface{}) (bool, error)
+	OneComplex(update interface{}) (bool, error)
 	Many(update interface{}) (int64, error)
 }
 
@@ -200,6 +229,22 @@ func (fo MongoUpdateOperation) One(update interface{}) (bool, error) {
 	return result.ModifiedCount == 1, nil
 }
 
+func (fo MongoUpdateOperation) OneComplex(update interface{}) (bool, error) {
+	if fo.err != nil {
+		return false, fo.err
+	}
+
+	result, err := fo.collection.UpdateOne(fo.ctx, fo.filter, update, fo.opt)
+	if err != nil {
+		return false, err
+	}
+
+	if fo.upsertResult != nil {
+		*(fo.upsertResult) = result.UpsertedID
+	}
+	return result.ModifiedCount == 1, nil
+}
+
 func (fo MongoUpdateOperation) Many(update interface{}) (int64, error) {
 	if fo.err != nil {
 		return 0, fo.err
@@ -238,8 +283,11 @@ func (storage *MongoStorage) Update(collectionName string) UpdateOperation {
 type FindOperation interface {
 	Context(ctx context.Context) FindOperation
 	Filter(filter OrderedDocument) FindOperation
+	Projection(filter OrderedDocument) FindOperation
 	Sort(field string, ascending bool) FindOperation
 	Limit(n int64) FindOperation
+	Skip(n int64) FindOperation
+	MaxTime(duration time.Duration) FindOperation
 	First(result interface{}) error
 	All(results interface{}) error
 }
@@ -247,6 +295,7 @@ type FindOperation interface {
 type MongoFindOperation struct {
 	collection *mongo.Collection
 	filter     OrderedDocument
+	projection OrderedDocument
 	ctx        context.Context
 	optFind    *options.FindOptions
 	optFindOne *options.FindOneOptions
@@ -266,11 +315,30 @@ func (fo MongoFindOperation) Filter(filter OrderedDocument) FindOperation {
 	return fo
 }
 
+func (fo MongoFindOperation) Projection(projection OrderedDocument) FindOperation {
+	for _, elem := range projection {
+		fo.projection = append(fo.projection, primitive.E{Key: elem.Key, Value: elem.Value})
+	}
+	fo.optFindOne.SetProjection(fo.projection)
+	fo.optFind.SetProjection(fo.projection)
+	return fo
+}
+
 func (fo MongoFindOperation) Limit(n int64) FindOperation {
 	fo.optFind.SetLimit(n)
 	return fo
 }
 
+func (fo MongoFindOperation) Skip(n int64) FindOperation {
+	fo.optFind.SetSkip(n)
+	return fo
+}
+
+func (fo MongoFindOperation) MaxTime(duration time.Duration) FindOperation {
+	fo.optFind.SetMaxTime(duration)
+	return fo
+}
+
 func (fo MongoFindOperation) Sort(field string, ascending bool) FindOperation {
 	var sort int
 	if ascending {
@@ -321,6 +389,7 @@ func (storage *MongoStorage) Find(collectionName string) FindOperation {
 	op := MongoFindOperation{
 		collection: collection,
 		filter:     OrderedDocument{},
+		projection: OrderedDocument{},
 		optFind:    options.Find(),
 		optFindOne: options.FindOne(),
 		sorts:      OrderedDocument{},