api/database/mongodb/db.go

package mongodb

import (
	"context"
	"fmt"
	"github.com/globalsign/mgo/bson"
	"time"

	"git.aiterp.net/rpdata/api/internal/config"
	"git.aiterp.net/rpdata/api/repositories"
	"github.com/globalsign/mgo"
)

type MongoDB struct {
	session *mgo.Session

	changes    repositories.ChangeRepository
	characters repositories.CharacterRepository
	tags       repositories.TagRepository
}

func (m *MongoDB) Changes() repositories.ChangeRepository {
	return m.changes
}

func (m *MongoDB) Characters() repositories.CharacterRepository {
	return m.characters
}

func (m *MongoDB) Tags() repositories.TagRepository {
	return m.tags
}

func (m *MongoDB) Close(ctx context.Context) error {
	m.session.Close()
	return nil
}

// Init initializes the mongodb database
func Init(cfg config.Database) (*MongoDB, error) {
	port := cfg.Port
	if port <= 0 {
		port = 27017
	}

	session, err := mgo.DialWithInfo(&mgo.DialInfo{
		Addrs:     []string{fmt.Sprintf("%s:%d", cfg.Host, port)},
		Timeout:   30 * time.Second,
		Database:  cfg.Db,
		Username:  cfg.Username,
		Password:  cfg.Password,
		Mechanism: cfg.Mechanism,
		Source:    cfg.Db,
	})
	if err != nil {
		return nil, err
	}

	db := session.DB(cfg.Db)

	characters, err := newCharacterRepository(db)
	if err != nil {
		session.Close()
		return nil, err
	}

	changes, err := newChangeRepository(db)
	if err != nil {
		session.Close()
		return nil, err
	}

	return &MongoDB{
		session: session,

		changes:    changes,
		characters: characters,
		tags:       newTagRepository(db),
	}, nil
}

type counter struct {
	coll     *mgo.Collection
	category string
	name     string
}

func newCounter(db *mgo.Database, category, name string) *counter {
	return &counter{
		coll:     db.C("core.counters"),
		category: category,
		name:     name,
	}
}

func (c *counter) WithName(name string) *counter {
	return &counter{
		coll:     c.coll,
		category: c.category,
		name:     name,
	}
}

func (c *counter) WithCategory(category string) *counter {
	return &counter{
		coll:     c.coll,
		category: category,
		name:     c.name,
	}
}

func (c *counter) With(category, name string) *counter {
	return &counter{
		coll:     c.coll,
		category: category,
		name:     name,
	}
}

func (c *counter) Increment(amount int) (int, error) {
	type counterDoc struct {
		ID    string `bson:"_id"`
		Value int    `bson:"value"`
	}

	id := c.category + "." + c.name
	doc := counterDoc{}

	_, err := c.coll.Find(bson.M{"_id": id}).Apply(mgo.Change{
		Update:    bson.M{"$inc": bson.M{"value": amount}},
		Upsert:    true,
		ReturnNew: true,
	}, &doc)
	if err != nil {
		return -1, err
	}

	return doc.Value, nil
}