From ee88d4bbfd2c25771c161a0edc685a59b427ddf8 Mon Sep 17 00:00:00 2001
From: Bel LaPointe <breel@qualtrics.com>
Date: Thu, 14 Mar 2019 08:56:39 -0600
Subject: [PATCH] add mongo

---
 db_test.go | 20 +++++++++++
 mongo.go   | 97 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 117 insertions(+)
 create mode 100644 mongo.go

diff --git a/db_test.go b/db_test.go
index e95467e..3962b72 100644
--- a/db_test.go
+++ b/db_test.go
@@ -72,6 +72,26 @@ func TestImplementations(t *testing.T) {
 		cases = append(cases, riak)
 	}
 
+	mongoLN, err := net.Listen("tcp", "localhost:27017")
+	if err == nil {
+		defer mongoLN.Close()
+		go func() {
+			for {
+				conn, err := mongoLN.Accept()
+				if err == nil {
+					conn.Close()
+				}
+			}
+		}()
+	}
+	if mongo1, err := NewMongo("localhost:27017"); err == nil {
+		cases = append(cases, mongo1)
+	} else if mongo2, err := NewMongo("localhost:27017", "root", "pass"); err == nil {
+		cases = append(cases, mongo2)
+	} else {
+		t.Errorf("cannot make mongo: %v", err)
+	}
+
 	validKey := "key"
 	validValue := []byte("value")
 
diff --git a/mongo.go b/mongo.go
new file mode 100644
index 0000000..6256e5d
--- /dev/null
+++ b/mongo.go
@@ -0,0 +1,97 @@
+package storage
+
+import (
+	"context"
+	"strings"
+	"time"
+
+	"github.com/mongodb/mongo-go-driver/mongo"
+	"github.com/mongodb/mongo-go-driver/mongo/options"
+	"github.com/mongodb/mongo-go-driver/mongo/readconcern"
+	"github.com/mongodb/mongo-go-driver/mongo/readpref"
+	"github.com/mongodb/mongo-go-driver/mongo/writeconcern"
+	"gopkg.in/mgo.v2/bson"
+)
+
+type Mongo struct {
+	db *mongo.Client
+}
+
+func NewMongo(addr string, auth ...string) (*Mongo, error) {
+	addr = strings.TrimPrefix(addr, "mongodb://")
+	var credentials *options.ClientOptions = nil
+	if len(auth) == 2 {
+		credentials = options.Client().SetAuth(options.Credential{
+			Username: auth[0],
+			Password: auth[1],
+		})
+	}
+	db, err := mongo.NewClientWithOptions(
+		"mongodb://"+addr,
+		options.Client().SetReadConcern(readconcern.Local()),
+		options.Client().SetReadPreference(readpref.PrimaryPreferred()),
+		options.Client().SetWriteConcern(writeconcern.New(
+			writeconcern.W(1),
+			writeconcern.J(false),
+		)),
+		credentials,
+	)
+	if err != nil {
+		return nil, err
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	if err := db.Connect(ctx); err != nil {
+		return nil, err
+	}
+	if err := db.Ping(ctx, nil); err != nil {
+		return nil, err
+	}
+	if _, err := db.ListDatabaseNames(ctx, &bson.M{}); err != nil {
+		return nil, err
+	}
+	return &Mongo{db: db}, nil
+}
+
+func (mg *Mongo) Get(key string) ([]byte, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	collection := mg.db.Database(DefaultNamespace).Collection(DefaultNamespace)
+	filter := bson.M{"_id": key}
+	cursor, err := collection.Find(ctx, filter)
+	if err != nil {
+		return nil, err
+	}
+	defer cursor.Close(ctx)
+	if !cursor.Next(ctx) {
+		return nil, ErrNotFound
+	}
+
+	elem, err := cursor.DecodeBytes()
+	if err != nil {
+		return nil, err
+	}
+	raw, err := elem.LookupErr("value")
+	if err != nil {
+		return nil, err
+	}
+	_, b, ok := raw.BinaryOK()
+	if !ok {
+		return nil, ErrNotImpl
+	}
+	return b, nil
+}
+
+func (mg *Mongo) Set(key string, value []byte) error {
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	collection := mg.db.Database(DefaultNamespace).Collection(DefaultNamespace)
+	filter := bson.M{"_id": key}
+	document := bson.M{"value": value}
+	_, err := collection.ReplaceOne(ctx, filter, document, options.Replace().SetUpsert(true))
+	return err
+}
+
+func (mg *Mongo) Close() error {
+	return nil
+}