Merge pull request #17 from Zondax/add/GetMongoDoc

Update on mongodb wrapper 
Minor refactor

Author: Ezequiel Raynaudo

connections/data_source.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	ds "github.com/Zondax/zindexer/connections/data_store"
 	"github.com/coinbase/rosetta-sdk-go/client"
+	"go.mongodb.org/mongo-driver/mongo"
 	"gorm.io/gorm"
 	"time"
 )
@@ -14,10 +15,11 @@ const (
 
 type DataSource struct {
 	// data sources
-	DbConn        *gorm.DB
-	RosettaClient *client.APIClient
-	NodeClient    interface{}
-	DataStore     ds.DataStoreClient
+	DatabasePostgres *gorm.DB
+	DatabaseMongo    *mongo.Client
+	RosettaClient    *client.APIClient
+	NodeClient       interface{}
+	DataStore        ds.DataStoreClient
 	// common
 	Ctx        context.Context
 	RetryDelay time.Duration
@@ -49,9 +51,15 @@ func WithRetryDelay(delay time.Duration) SourceOption {
 	}
 }
 
-func WithDBConnection(dbConn *gorm.DB) SourceOption {
+func WithPostgresDB(dbConn *gorm.DB) SourceOption {
 	return func(w *DataSource) {
-		w.DbConn = dbConn
+		w.DatabasePostgres = dbConn
+	}
+}
+
+func WithMongoDB(dbConn *mongo.Client) SourceOption {
+	return func(w *DataSource) {
+		w.DatabaseMongo = dbConn
 	}
 }
 
@@ -69,10 +77,10 @@ func WithNodeClient(node interface{}) SourceOption {
 
 func WithDataStore(cfg ds.DataStoreConfig) SourceOption {
 	return func(w *DataSource) {
-		client, err := ds.NewDataStoreClient(cfg)
+		storeClient, err := ds.NewDataStoreClient(cfg)
 		if err != nil {
 			panic(err)
 		}
-		w.DataStore = client
+		w.DataStore = storeClient
 	}
 }

connections/database/graphql.go connections/database/graphql/graphql.goconnections/database/graphql.go renamed to connections/database/graphql/graphql.go

@@ -1,8 +1,9 @@
-package database
+package graphql
 
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/Zondax/zindexer/connections/database"
 	"github.com/hasura/go-graphql-client"
 	"go.uber.org/zap"
 	"net/http"
@@ -28,7 +29,7 @@ type GraphqlSubscriptionClient struct {
 func NewGraphqlQueryClient(host string, token string) GraphqlClient {
 	transport := http.DefaultTransport
 	if token != "" {
-		transport = AuthHeaderTransport{Transport: http.DefaultTransport, Token: token}
+		transport = database.AuthHeaderTransport{Transport: http.DefaultTransport, Token: token}
 	}
 
 	customHttpClient := http.Client{Transport: transport}

connections/database/mongoDB.go

@@ -0,0 +1,71 @@
+package mongodb
+
+import (
+	"context"
+	"fmt"
+	"github.com/Zondax/zindexer/connections/database"
+	"go.uber.org/zap"
+	"time"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
+	"go.mongodb.org/mongo-driver/mongo/readpref"
+)
+
+type MongoConnection struct {
+	db *mongo.Client
+}
+
+func NewMongoConnection(params *database.DBConnectionParams) (*MongoConnection, error) {
+	uri := params.URI
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	client, err := mongo.Connect(ctx, options.Client().ApplyURI(uri))
+	if err != nil {
+		return nil, err
+	}
+
+	// Ping the primary
+	if err := client.Ping(ctx, readpref.Primary()); err != nil {
+		return nil, err
+	}
+
+	return &MongoConnection{db: client}, nil
+}
+
+func (c *MongoConnection) GetDB() *mongo.Client {
+	return c.db
+}
+
+func Connect(params *database.DBConnectionParams) (*mongo.Client, error) {
+	conn, err := NewMongoConnection(params)
+	if err != nil {
+		return nil, err
+	}
+
+	return conn.GetDB(), nil
+}
+
+func (c *MongoConnection) GetMongoDoc(collection *mongo.Collection, docId string) (bson.M, error) {
+	zap.S().Debug("document with id:%v \n", docId)
+	opts := options.FindOne()
+	var result bson.M
+	readErr := collection.FindOne(
+		context.TODO(),
+		bson.D{{Key: "_id", Value: docId}},
+		opts,
+	).Decode(&result)
+
+	if readErr != nil {
+		// ErrNoDocuments means that the filter did not match any documents in
+		// the collection.
+		if readErr == mongo.ErrNoDocuments {
+			return nil, fmt.Errorf("no document found")
+		}
+		return nil, readErr
+	}
+
+	return result, nil
+}

connections/database/mongoDB_test.go

@@ -0,0 +1,39 @@
+package mongodb
+
+import (
+	"context"
+	"crypto/rand"
+	"fmt"
+	"github.com/Zondax/zindexer/connections/database"
+	"go.mongodb.org/mongo-driver/bson"
+	"os"
+	"testing"
+)
+
+func TestBuffer_InsertRead(t *testing.T) {
+	params := &database.DBConnectionParams{URI: os.Getenv("MONGO_URI")}
+	client, err := NewMongoConnection(params)
+	if err != nil {
+		fmt.Println(err)
+		t.Fail()
+	}
+
+	coll := client.GetDB().Database("test").Collection("sample")
+	key, _ := rand.Prime(rand.Reader, 32)
+	res, err := coll.InsertOne(context.TODO(), bson.D{{Key: "name", Value: "Alice"}, {Key: "_id", Value: key.String()}})
+	if err != nil {
+		fmt.Println(err)
+		t.Fail()
+	}
+	fmt.Printf("inserted document with ID %v\n", res.InsertedID)
+
+	// Get the inserted file
+	result, err := client.GetMongoDoc(coll, key.String())
+	if err != nil {
+		fmt.Printf("Failed to get with error %v", err)
+		t.Fail()
+		return
+	}
+
+	fmt.Printf("found document %v\n", result)
+}

connections/database/mongodb/mongoDB.go

@@ -1,7 +1,8 @@
-package database
+package postgres
 
 import (
 	"fmt"
+	"github.com/Zondax/zindexer/connections/database"
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 )
@@ -14,7 +15,7 @@ type DBConnectionConfig struct {
 	Gorm *gorm.Config
 }
 
-func NewPostgresConnection(params *DBConnectionParams, config DBConnectionConfig) (*GormConnection, error) {
+func NewPostgresConnection(params *database.DBConnectionParams, config DBConnectionConfig) (*GormConnection, error) {
 	dsn, err := params.GetDSN()
 	if err != nil {
 		return nil, fmt.Errorf("failed to retrieve dsn")
@@ -24,22 +25,20 @@ func NewPostgresConnection(params *DBConnectionParams, config DBConnectionConfig
 	if config.Gorm != nil {
 		dbConfig = *config.Gorm
 	}
-	conn, err := gorm.Open(postgres.Open(dsn), &dbConfig)
+	c, err := gorm.Open(postgres.Open(dsn), &dbConfig)
 
 	if err != nil {
 		return nil, fmt.Errorf("failed to dial connect to db '%s@%s:%s': %v", params.Name, params.Host, params.Port, err)
 	}
 
-	return &GormConnection{
-		db: conn,
-	}, nil
+	return &GormConnection{db: c}, nil
 }
 
 func (c *GormConnection) GetDB() *gorm.DB {
 	return c.db
 }
 
-func ConnectDB(params DBConnectionParams, config DBConnectionConfig) (*gorm.DB, error) {
+func Connect(params database.DBConnectionParams, config DBConnectionConfig) (*gorm.DB, error) {
 	conn, err := NewPostgresConnection(&params, config)
 	if err != nil {
 		return nil, err

connections/database/mongodb/mongoDB_test.go

@@ -3,8 +3,6 @@ package database
 import (
 	"fmt"
 	"net/http"
-
-	"gorm.io/gorm"
 )
 
 type AuthHeaderTransport struct {
@@ -17,19 +15,6 @@ func (adt AuthHeaderTransport) RoundTrip(req *http.Request) (*http.Response, err
 	return adt.Transport.RoundTrip(req)
 }
 
-type DBConnection interface {
-	GetDB() *gorm.DB
-}
-
-type IDBQueryClient interface {
-	GetDB() interface{}
-	Connect() error
-}
-
-type DBQueryClient struct {
-	Client IDBQueryClient
-}
-
 type DBConnectionParams struct {
 	User     string
 	Password string