WilliesSchematicsWorld/schematic-service-go/dao/mongodao.go

package dao

import (
	"context"
	"crypto/md5"
	"fmt"
	"io"
	"log"
	"strings"
	"time"

	"github.com/willie68/schematic-service-go/config"
	slicesutils "github.com/willie68/schematic-service-go/internal"
	"github.com/willie68/schematic-service-go/model"

	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/gridfs"
	"go.mongodb.org/mongo-driver/mongo/options"
)

const timeout = 1 * time.Minute
const attachmentsCollectionName = "attachments"
const schematicsCollectionName = "schematics"
const tagsCollectionName = "tags"
const manufacturersCollectionName = "manufacturers"
const usersCollectionName = "users"

var client *mongo.Client
var mongoConfig config.MongoDB
var bucket gridfs.Bucket
var database mongo.Database
var tags []string
var manufacturers []string
var users map[string]string

// InitDB initialise the mongodb connection, build up all collections and indexes
func InitDB(MongoConfig config.MongoDB) {
	mongoConfig = MongoConfig
	//	uri := fmt.Sprintf("mongodb://%s:%s@%s:%d", mongoConfig.Username, mongoConfig.Password, mongoConfig.Host, mongoConfig.Port)
	uri := fmt.Sprintf("mongodb://%s:%d", mongoConfig.Host, mongoConfig.Port)
	clientOptions := options.Client()
	clientOptions.ApplyURI(uri)
	clientOptions.Auth = &options.Credential{Username: mongoConfig.Username, Password: mongoConfig.Password, AuthSource: mongoConfig.AuthDB}
	var err error
	client, err = mongo.NewClient(clientOptions)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
	}
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	err = client.Connect(ctx)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
	}
	database = *client.Database(mongoConfig.Database)

	myBucket, err := gridfs.NewBucket(&database, options.GridFSBucket().SetName(attachmentsCollectionName))
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
	}
	bucket = *myBucket

	initIndexSchematics()
	initIndexTags()
	initIndexManufacturers()

	tags = make([]string, 0)
	manufacturers = make([]string, 0)
	users = make(map[string]string)
	initTags()
	initManufacturers()
	initUsers()
}

func initIndexSchematics() {
	collection := database.Collection(schematicsCollectionName)
	indexView := collection.Indexes()
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	cursor, err := indexView.List(ctx)
	if err != nil {
		log.Fatal(err)
	}
	defer cursor.Close(ctx)
	myIndexes := make([]string, 0)
	for cursor.Next(ctx) {
		var index bson.M
		if err = cursor.Decode(&index); err != nil {
			log.Fatal(err)
		}
		myIndexes = append(myIndexes, index["name"].(string))
	}

	for _, name := range myIndexes {
		log.Println(name)
	}
	if !slicesutils.Contains(myIndexes, "manufaturer") {
		ctx, _ = context.WithTimeout(context.Background(), timeout)
		models := []mongo.IndexModel{
			{
				Keys:    bson.D{{"manufacturer", 1}},
				Options: options.Index().SetName("manufacturer").SetCollation(&options.Collation{Locale: "en", Strength: 2}),
			},
			{
				Keys:    bson.D{{"model", 1}},
				Options: options.Index().SetName("model").SetCollation(&options.Collation{Locale: "en", Strength: 2}),
			},
			{
				Keys:    bson.D{{"tags", 1}},
				Options: options.Index().SetName("tags").SetCollation(&options.Collation{Locale: "en", Strength: 2}),
			},
			{
				Keys:    bson.D{{"subtitle", 1}},
				Options: options.Index().SetName("subtitle").SetCollation(&options.Collation{Locale: "en", Strength: 2}),
			},
			{
				Keys:    bson.D{{"manufacturer", "text"}, {"model", "text"}, {"tags", "text"}, {"subtitle", "text"}, {"description", "text"}, {"owner", "text"}},
				Options: options.Index().SetName("$text"),
			},
		}

		// Specify the MaxTime option to limit the amount of time the operation can run on the server
		opts := options.CreateIndexes().SetMaxTime(2 * time.Second)
		names, err := indexView.CreateMany(context.TODO(), models, opts)
		if err != nil {
			log.Fatal(err)
		}
		log.Print("create indexes:")
		for _, name := range names {
			log.Println(name)
		}
	}
}

func initIndexTags() {
	collection := database.Collection(tagsCollectionName)
	indexView := collection.Indexes()
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	cursor, err := indexView.List(ctx)
	if err != nil {
		log.Fatal(err)
	}
	defer cursor.Close(ctx)
	myIndexes := make([]string, 0)
	for cursor.Next(ctx) {
		var index bson.M
		if err = cursor.Decode(&index); err != nil {
			log.Fatal(err)
		}
		myIndexes = append(myIndexes, index["name"].(string))
	}

	for _, name := range myIndexes {
		log.Println(name)
	}
	if !slicesutils.Contains(myIndexes, "name") {
		ctx, _ = context.WithTimeout(context.Background(), timeout)
		models := []mongo.IndexModel{
			{
				Keys:    bson.D{{"name", 1}},
				Options: options.Index().SetUnique(true).SetName("name").SetCollation(&options.Collation{Locale: "en", Strength: 2}),
			},
		}

		// Specify the MaxTime option to limit the amount of time the operation can run on the server
		opts := options.CreateIndexes().SetMaxTime(2 * time.Second)
		names, err := indexView.CreateMany(context.TODO(), models, opts)
		if err != nil {
			log.Fatal(err)
		}
		log.Print("create indexes:")
		for _, name := range names {
			log.Println(name)
		}
	}
}

func initIndexManufacturers() {
	collection := database.Collection(manufacturersCollectionName)
	indexView := collection.Indexes()
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	cursor, err := indexView.List(ctx)
	if err != nil {
		log.Fatal(err)
	}
	defer cursor.Close(ctx)
	myIndexes := make([]string, 0)
	for cursor.Next(ctx) {
		var index bson.M
		if err = cursor.Decode(&index); err != nil {
			log.Fatal(err)
		}
		myIndexes = append(myIndexes, index["name"].(string))
	}

	for _, name := range myIndexes {
		log.Println(name)
	}
	if !slicesutils.Contains(myIndexes, "name") {
		ctx, _ = context.WithTimeout(context.Background(), timeout)
		models := []mongo.IndexModel{
			{
				Keys:    bson.D{{"name", 1}},
				Options: options.Index().SetUnique(true).SetName("name").SetCollation(&options.Collation{Locale: "en", Strength: 2}),
			},
		}

		// Specify the MaxTime option to limit the amount of time the operation can run on the server
		opts := options.CreateIndexes().SetMaxTime(2 * time.Second)
		names, err := indexView.CreateMany(context.TODO(), models, opts)
		if err != nil {
			log.Fatal(err)
		}
		log.Print("create indexes:")
		for _, name := range names {
			log.Println(name)
		}
	}
}

func initTags() {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	tagsCollection := database.Collection(tagsCollectionName)
	cursor, err := tagsCollection.Find(ctx, bson.M{})
	if err != nil {
		log.Fatal(err)
	}
	defer cursor.Close(ctx)
	for cursor.Next(ctx) {
		var tag bson.M
		if err = cursor.Decode(&tag); err != nil {
			log.Fatal(err)
		} else {
			tags = append(tags, tag["name"].(string))
		}
	}
}

func initManufacturers() {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	manufacturersCollection := database.Collection(manufacturersCollectionName)
	cursor, err := manufacturersCollection.Find(ctx, bson.M{})
	if err != nil {
		log.Fatal(err)
	}
	defer cursor.Close(ctx)
	for cursor.Next(ctx) {
		var manufacturer bson.M
		if err = cursor.Decode(&manufacturer); err != nil {
			log.Fatal(err)
		} else {
			manufacturers = append(manufacturers, manufacturer["name"].(string))
		}
	}
}

func initUsers() {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	usersCollection := database.Collection(usersCollectionName)
	cursor, err := usersCollection.Find(ctx, bson.M{})
	if err != nil {
		log.Fatal(err)
	}
	defer cursor.Close(ctx)
	for cursor.Next(ctx) {
		var user bson.M
		if err = cursor.Decode(&user); err != nil {
			log.Fatal(err)
		} else {
			username := user["name"].(string)
			password := user["password"].(string)
			if !strings.HasPrefix(password, "md5:") {
				hash := md5.Sum([]byte(password))
				password = fmt.Sprintf("md5:%x", hash)
			}
			users[username] = password
		}
	}
}

// AddFile adding a file to the storage, stream like
func AddFile(filename string, reader io.Reader) (string, error) {
	uploadOpts := options.GridFSUpload().SetMetadata(bson.D{{"tag", "tag"}})

	fileID, err := bucket.UploadFromStream(filename, reader, uploadOpts)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return "", err
	}
	log.Printf("Write file to DB was successful. File id: %s \n", fileID)
	id := fileID.Hex()
	return id, nil
}

// CreateSchematic creating a new schematic in the database
func CreateSchematic(schematic model.Schematic) (string, error) {

	for _, tag := range schematic.Tags {
		if !slicesutils.Contains(tags, tag) {
			CreateTag(tag)
		}
	}

	if !slicesutils.Contains(manufacturers, schematic.Manufacturer) {
		CreateManufacturer(schematic.Manufacturer)
	}

	ctx, _ := context.WithTimeout(context.Background(), timeout)
	collection := database.Collection(schematicsCollectionName)
	result, err := collection.InsertOne(ctx, schematic)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return "", err
	}
	filter := bson.M{"_id": result.InsertedID}
	err = collection.FindOne(ctx, filter).Decode(&schematic)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return "", err
	}
	switch v := result.InsertedID.(type) {
	case primitive.ObjectID:
		return v.Hex(), nil
	}
	return "", nil
}

// GetSchematic getting a sdingle schematic
func GetSchematic(schematicID string) (model.Schematic, error) {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	schematicCollection := database.Collection(schematicsCollectionName)
	objectId, _ := primitive.ObjectIDFromHex(schematicID)
	result := schematicCollection.FindOne(ctx, bson.M{"_id": objectId})
	var schematic model.Schematic
	if err := result.Decode(&schematic); err != nil {
		log.Print(err)
		return model.Schematic{}, err
	} else {
		return schematic, nil
	}
}

func GetFile(fileid string, stream io.Writer) error {
	objectID, err := primitive.ObjectIDFromHex(fileid)
	if err != nil {
		log.Print(err)
		return err
	}
	dStream, err := bucket.DownloadToStream(objectID, stream)
	if err != nil {
		log.Print(err)
		return err
	}
	fmt.Printf("File size to download: %v \n", dStream)
	return nil
}

// GetSchematics getting a sdingle schematic
func GetSchematics(query string, offset int, limit int, owner string) ([]model.Schematic, error) {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	schematicCollection := database.Collection(schematicsCollectionName)
	queryDoc := bson.M{}
	err := bson.UnmarshalExtJSON([]byte(query), false, queryDoc)
	if err != nil {
		log.Print(err)
		return nil, err
	}
	cursor, err := schematicCollection.Find(ctx, queryDoc, &options.FindOptions{Collation: &options.Collation{Locale: "en", Strength: 2}})
	if err != nil {
		log.Print(err)
		return nil, err
	}
	defer cursor.Close(ctx)
	schematics := make([]model.Schematic, 0)
	count := 0
	docs := 0
	for cursor.Next(ctx) {
		if count >= offset {
			if docs < limit {
				var schematic model.Schematic
				if err = cursor.Decode(&schematic); err != nil {
					log.Print(err)
					return nil, err
				} else {
					if !schematic.PrivateFile || schematic.Owner == owner {
						schematics = append(schematics, schematic)
						docs++
					}
				}
			} else {
				break
			}
		}
		count++
	}
	return schematics, nil
}

// CreateTag create a new tag in the storage
func CreateTag(tag string) error {
	tag = strings.ToLower(tag)
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	collection := database.Collection(tagsCollectionName)
	tagModel := bson.M{"name": tag}
	_, err := collection.InsertOne(ctx, tagModel)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return err
	}
	tags = append(tags, tag)
	return nil
}

// CreateManufacturer create a new manufacturer in the storage
func CreateManufacturer(manufacturer string) error {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	collection := database.Collection(manufacturersCollectionName)
	manufacturerModel := bson.M{"name": manufacturer}
	_, err := collection.InsertOne(ctx, manufacturerModel)
	if err != nil {
		fmt.Printf("error: %s\n", err.Error())
		return err
	}
	manufacturers = append(manufacturers, manufacturer)
	return nil
}

func GetTags() []string {
	return tags
}

func GetManufacturers() []string {
	return manufacturers
}

func GetTagsCount() int {
	return len(tags)
}

func GetManufacturersCount() int {
	return len(manufacturers)
}

func CheckUser(username string, password string) bool {
	pwd, ok := users[username]
	if ok {
		if pwd == password {
			return true
		}
	}
	return false
}

func DropAll() {
	ctx, _ := context.WithTimeout(context.Background(), timeout)
	collectionNames, err := database.ListCollectionNames(ctx, bson.D{}, &options.ListCollectionsOptions{})
	if err != nil {
		log.Fatal(err)
	}
	for _, name := range collectionNames {
		collection := database.Collection(name)
		err = collection.Drop(ctx)
		if err != nil {
			log.Fatal(err)
		}
	}
}