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"
)

var DisableFixes bool

type MongoDB struct {
	session *mgo.Session

	changes    repositories.ChangeRepository
	channels   *channelRepository
	characters repositories.CharacterRepository
	tags       repositories.TagRepository
	logs       *logRepository
	posts      *postRepository
	files      *fileRepository
	stories    repositories.StoryRepository
	chapters   repositories.ChapterRepository
	comments   repositories.CommentRepository
	keys       repositories.KeyRepository
	users      repositories.UserRepository
}

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) Logs() repositories.LogRepository {
	return m.logs
}

func (m *MongoDB) Posts() repositories.PostRepository {
	return m.posts
}

func (m *MongoDB) Channels() repositories.ChannelRepository {
	return m.channels
}

func (m *MongoDB) Files() repositories.FileRepository {
	return m.files
}

func (m *MongoDB) Stories() repositories.StoryRepository {
	return m.stories
}

func (m *MongoDB) Chapters() repositories.ChapterRepository {
	return m.chapters
}

func (m *MongoDB) Comments() repositories.CommentRepository {
	return m.comments
}

func (m *MongoDB) Keys() repositories.KeyRepository {
	return m.keys
}

func (m *MongoDB) Users() repositories.UserRepository {
	return m.users
}

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, cfg.RestoreIDs)
	if err != nil {
		session.Close()
		return nil, err
	}

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

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

	logs, err := newLogRepository(db, cfg.RestoreIDs)
	if err != nil {
		session.Close()
		return nil, err
	}

	posts, err := newPostRepository(db, cfg.RestoreIDs)
	if err != nil {
		session.Close()
		return nil, err
	}

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

	stories, err := newStoryRepository(db, cfg.RestoreIDs)
	if err != nil {
		session.Close()
		return nil, err
	}

	chapters, err := newChapterRepository(db, cfg.RestoreIDs)
	if err != nil {
		session.Close()
		return nil, err
	}

	comments, err := newCommentRepository(db, cfg.RestoreIDs)
	if err != nil {
		session.Close()
		return nil, err
	}

	keys, err := newKeyRepository(db, cfg.RestoreIDs)
	if err != nil {
		session.Close()
		return nil, err
	}

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

	if !DisableFixes {
		go posts.fixPositions(logs)
	}

	return &MongoDB{
		session: session,

		changes:    changes,
		characters: characters,
		channels:   channels,
		tags:       newTagRepository(db),
		stories:    stories,
		chapters:   chapters,
		comments:   comments,
		logs:       logs,
		posts:      posts,
		files:      files,
		keys:       keys,
		users:      users,
	}, 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) Bump(amount int) error {
	id := c.category + "." + c.name

	err := c.coll.Update(bson.M{
		"_id":   id,
		"value": bson.M{"$lt": amount},
	}, bson.M{
		"value": amount,
	})
	if err != nil {
		if err == mgo.ErrNotFound {
			return c.coll.Insert(bson.M{
				"_id":   id,
				"value": amount,
			})
		}

		return err
	}

	return nil
}

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
}