trantor/lib/database/mgo.go

package database

import (
	"errors"

	mgo "gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

const (
	visited_coll    = "visited"
	downloaded_coll = "downloaded"
	tags_coll       = "tags"
)

type mgoDB struct {
	session *mgo.Session
	name    string
}

func (db *mgoDB) initIndexes() {
	dbCopy := db.session.Copy()
	booksColl := dbCopy.DB(db.name).C(books_coll)
	go indexBooks(booksColl)
	statsColl := dbCopy.DB(db.name).C(stats_coll)
	go indexStats(statsColl)
	newsColl := dbCopy.DB(db.name).C(news_coll)
	go indexNews(newsColl)
}

func (db *mgoDB) Close() {
	db.session.Close()
}

func (db *mgoDB) Copy() DB {
	dbCopy := new(mgoDB)
	dbCopy.session = db.session.Copy()
	dbCopy.name = db.name
	return dbCopy
}

func (db *mgoDB) AddBook(book map[string]interface{}) error {
	booksColl := db.session.DB(db.name).C(books_coll)
	return addBook(booksColl, book)
}

func (db *mgoDB) GetBooks(query string, length int, start int) (books []Book, num int, err error) {
	booksColl := db.session.DB(db.name).C(books_coll)
	return getBooks(booksColl, query, length, start)
}

func (db *mgoDB) GetBooksIter() Iter {
	booksColl := db.session.DB(db.name).C(books_coll)
	return getBooksIter(booksColl)
}

func (db *mgoDB) GetNewBooks(query string, length int, start int) (books []Book, num int, err error) {
	booksColl := db.session.DB(db.name).C(books_coll)
	return getNewBooks(booksColl, query, length, start)
}

func (db *mgoDB) GetBookId(id string) (Book, error) {
	booksColl := db.session.DB(db.name).C(books_coll)
	return getBookId(booksColl, id)
}

func (db *mgoDB) DeleteBook(id string) error {
	booksColl := db.session.DB(db.name).C(books_coll)
	return deleteBook(booksColl, id)
}

func (db *mgoDB) UpdateBook(id string, data map[string]interface{}) error {
	booksColl := db.session.DB(db.name).C(books_coll)
	return updateBook(booksColl, id, data)
}

func (db *mgoDB) FlagBadQuality(id string, user string) error {
	booksColl := db.session.DB(db.name).C(books_coll)
	return flagBadQuality(booksColl, id, user)
}

func (db *mgoDB) ActiveBook(id string) error {
	booksColl := db.session.DB(db.name).C(books_coll)
	return activeBook(booksColl, id)
}

func (db *mgoDB) IsBookActive(id string) bool {
	booksColl := db.session.DB(db.name).C(books_coll)
	return isBookActive(booksColl, id)
}

func (db *mgoDB) User(name string) *User {
	userColl := db.session.DB(db.name).C(user_coll)
	return getUser(userColl, name)
}

func (db *mgoDB) AddUser(name string, pass string) error {
	userColl := db.session.DB(db.name).C(user_coll)
	return addUser(userColl, name, pass)
}

func (db *mgoDB) AddNews(text string) error {
	newsColl := db.session.DB(db.name).C(news_coll)
	return addNews(newsColl, text)
}

func (db *mgoDB) GetNews(num int, days int) (news []News, err error) {
	newsColl := db.session.DB(db.name).C(news_coll)
	return getNews(newsColl, num, days)
}

// TODO: split code in files
func (db *mgoDB) AddStats(stats interface{}) error {
	statsColl := db.session.DB(db.name).C(stats_coll)
	return statsColl.Insert(stats)
}

/* Get the most visited books
 */
func (db *mgoDB) GetVisitedBooks() (books []Book, err error) {
	visitedColl := db.session.DB(db.name).C(visited_coll)
	bookId, err := GetBooksVisited(visitedColl)
	if err != nil {
		return nil, err
	}

	books = make([]Book, len(bookId))
	for i, id := range bookId {
		booksColl := db.session.DB(db.name).C(books_coll)
		booksColl.Find(bson.M{"_id": id}).One(&books[i])
		books[i].Id = bson.ObjectId(books[i].Id).Hex()
	}
	return
}

func (db *mgoDB) UpdateMostVisited() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(visited_coll)
	return u.UpdateMostBooks("book")
}

/* Get the most downloaded books
 */
func (db *mgoDB) GetDownloadedBooks() (books []Book, err error) {
	downloadedColl := db.session.DB(db.name).C(downloaded_coll)
	bookId, err := GetBooksVisited(downloadedColl)
	if err != nil {
		return nil, err
	}

	books = make([]Book, len(bookId))
	for i, id := range bookId {
		booksColl := db.session.DB(db.name).C(books_coll)
		booksColl.Find(bson.M{"_id": id}).One(&books[i])
		books[i].Id = bson.ObjectId(books[i].Id).Hex()
	}
	return
}

func (db *mgoDB) UpdateDownloadedBooks() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(downloaded_coll)
	return u.UpdateMostBooks("download")
}

func (db *mgoDB) GetTags() ([]string, error) {
	tagsColl := db.session.DB(db.name).C(tags_coll)
	return GetTags(tagsColl)
}

func (db *mgoDB) UpdateTags() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(books_coll)
	u.dst = db.session.DB(db.name).C(tags_coll)
	return u.UpdateTags()
}

func (db *mgoDB) GetVisits(visitType VisitType) ([]Visits, error) {
	var coll *mgo.Collection
	switch visitType {
	case Hourly_visits:
		coll = db.session.DB(db.name).C(hourly_visits_coll)
	case Daily_visits:
		coll = db.session.DB(db.name).C(daily_visits_coll)
	case Monthly_visits:
		coll = db.session.DB(db.name).C(monthly_visits_coll)
	case Hourly_downloads:
		coll = db.session.DB(db.name).C(hourly_downloads_coll)
	case Daily_downloads:
		coll = db.session.DB(db.name).C(daily_downloads_coll)
	case Monthly_downloads:
		coll = db.session.DB(db.name).C(monthly_downloads_coll)
	default:
		return nil, errors.New("Not valid VisitType")
	}
	return GetVisits(coll)
}

func (db *mgoDB) UpdateHourVisits() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(hourly_visits_coll)
	return u.UpdateHourVisits(false)
}

func (db *mgoDB) UpdateDayVisits() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(daily_visits_coll)
	return u.UpdateDayVisits(false)
}

func (db *mgoDB) UpdateMonthVisits() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(monthly_visits_coll)
	return u.UpdateMonthVisits(false)
}

func (db *mgoDB) UpdateHourDownloads() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(hourly_downloads_coll)
	return u.UpdateHourVisits(true)
}

func (db *mgoDB) UpdateDayDownloads() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(daily_downloads_coll)
	return u.UpdateDayVisits(true)
}

func (db *mgoDB) UpdateMonthDownloads() error {
	var u dbUpdate
	u.src = db.session.DB(db.name).C(stats_coll)
	u.dst = db.session.DB(db.name).C(monthly_downloads_coll)
	return u.UpdateMonthVisits(true)
}