trantor/lib/database/books.go

package database

import (
	log "github.com/cihub/seelog"

	"strings"
	"time"

	"github.com/go-pg/pg/v10"
)

// Book metadata
type Book struct {
	ID          string `pg:"type:varchar(16)"`
	Title       string
	Authors     []string `pg:"authors,array"`
	Contributor string
	Publisher   string
	Description string
	Tags        []string `pg:"tags,array"`
	Date        string
	Lang        string `pg:"type:varchar(3)"`
	Isbn        string `pg:"type:varchar(13)"`
	FileSize    int    `pg:"type:integer"`
	FileHash    []byte
	Cover       bool      `pg:",use_zero"`
	Active      bool      `pg:",use_zero"`
	UploadDate  time.Time `pg:"type:timestamp"`
	Tsv         string    `pg:"type:tsvector"`
}

// AddBook to the database
func (db *pgDB) AddBook(book Book) error {
	emptyTime := time.Time{}
	if book.UploadDate == emptyTime {
		book.UploadDate = time.Now()
	}

	_, err := db.sql.Model(&book).Insert()
	return err
}

// GetBooks matching query
func (db *pgDB) GetBooks(query string, length int, start int) (books []Book, num int, err error) {
	return db.getBooks(true, query, length, start)
}

// GetNewBooks returns a list of books in the incoming queue and the number of books
// in the queue
func (db *pgDB) GetNewBooks(query string, length int, start int) (books []Book, num int, err error) {
	return db.getBooks(false, query, length, start)
}

func (db *pgDB) getBooks(active bool, query string, length int, start int) (books []Book, num int, err error) {
	rank := []string{}
	rankParams := []interface{}{}
	searchCondition := ""
	if active {
		searchCondition += "active is true"
	} else {
		searchCondition += "active is not true"
	}
	searchParams := []interface{}{}

	textQuery, columnQuerys, trigramQuerys := buildQuery(query)
	for _, c := range columnQuerys {
		searchCondition += " AND " + c.column + " = ?"
		searchParams = append(searchParams, c.value)
	}
	for _, c := range trigramQuerys {
		rank = append(rank, "word_similarity(?, "+c.column+")")
		rankParams = append(rankParams, c.value)
		searchCondition += " AND " + c.column + " %> ?"
		searchParams = append(searchParams, c.value)
	}
	if textQuery != "" {
		rank = append(rank, "ts_rank(tsv, to_tsquery_multilingual(?))")
		rankParams = append(rankParams, textQuery)
		searchCondition += " AND to_tsquery_multilingual(?) @@ tsv"
		searchParams = append(searchParams, textQuery)
	}

	order := "upload_date DESC"
	if len(rank) > 0 {
		order = strings.Join(rank, "+") + " DESC, upload_date DESC"
	}

	q := db.sql.Model(&books).
		Where(searchCondition, searchParams...).
		OrderExpr(order, rankParams...).
		Offset(start).
		Limit(length)
	if active {
		num, err = q.SelectAndCountEstimate(1000)
	} else {
		num, err = q.SelectAndCount()
	}
	return books, num, err
}

// GetBookID returns a the book with the specified id
func (db *pgDB) GetBookID(id string) (Book, error) {
	var book Book
	err := db.sql.Model(&book).
		Where("id = ?", id).
		Select()
	return book, err
}

// DeleteBook removes the book with id from the database
func (db *pgDB) DeleteBook(id string) error {
	_, err := db.sql.Model(&Book{}).
		Where("id = ?", id).
		Delete()
	return err
}

// UpdateBook metadata
func (db *pgDB) UpdateBook(id string, data map[string]interface{}) error {
	setCondition := ""
	params := []interface{}{}
	for col, val := range data {
		colValid := false
		for _, name := range []string{"title", "authors", "contributor", "publisher",
			"description", "tags", "date", "lang", "isbn"} {
			if col == name {
				colValid = true
				break
			}
		}
		if !colValid {
			continue
		}

		if len(setCondition) != 0 {
			setCondition += ", "
		}
		setCondition += col + " = ?"
		if col == "authors" || col == "tags" {
			params = append(params, pg.Array(val))
		} else {
			params = append(params, val)
		}
	}
	_, err := db.sql.Model(&Book{}).
		Set(setCondition, params...).
		Where("id = ?", id).
		Update()
	return err
}

// ActiveBook activates the book
func (db *pgDB) ActiveBook(id string) error {
	uploadDate := time.Now()
	_, err := db.sql.Model(&Book{}).
		Set("active = true, upload_date = ? ", uploadDate).
		Where("id = ?", id).
		Update()
	return err
}

// IsBookActive checks if the book is active
func (db *pgDB) IsBookActive(id string) bool {
	var active []bool
	err := db.sql.Model(&Book{}).
		Column("active").
		Where("id = ?", id).
		Select(&active)
	if err != nil || len(active) != 1 {
		return false
	}
	return active[0]
}

// ExistsBookHash checks if the given hash matches the hash of any book in the lbirary
func (db *pgDB) ExistsBookHash(hash []byte) bool {
	count, err := db.sql.Model(&Book{}).
		Where("file_hash = ?", hash).
		Count()
	if err != nil {
		log.Warnf("There was an error looking for the hash: %v", err)
		return false
	}
	return count > 0
}

type columnq struct {
	column string
	value  string
}

func buildQuery(query string) (textQuery string, columnQuerys []columnq, trigramQuerys []columnq) {
	tokens := extractTokens(query)
	for _, token := range tokens {
		if token == "" {
			continue
		}
		tag := strings.SplitN(token, ":", 2)
		value := ""
		if len(tag) > 1 {
			value = strings.Replace(tag[1], "%", "\\%", 0)
			value = strings.Replace(value, "_", "\\_", 0)
		}

		switch tag[0] {
		case "lang":
			columnQuerys = append(columnQuerys, columnq{"lang", value})
		case "isbn":
			columnQuerys = append(columnQuerys, columnq{"isbn", value})
		case "author":
			trigramQuerys = append(trigramQuerys, columnq{"text(authors)", value})
		case "title":
			trigramQuerys = append(trigramQuerys, columnq{"title", value})
		case "contributor":
			trigramQuerys = append(trigramQuerys, columnq{"contributor", value})
		case "publisher":
			trigramQuerys = append(trigramQuerys, columnq{"publisher", value})
		case "subject":
			trigramQuerys = append(trigramQuerys, columnq{"text(tags)", value})
		case "tag":
			trigramQuerys = append(trigramQuerys, columnq{"text(tags)", value})
		case "date":
			trigramQuerys = append(trigramQuerys, columnq{"date", value})
		default:
			if len(textQuery) != 0 {
				lastChar := textQuery[len(textQuery)-1:]
				if token[:1] != "&" && token[:1] != "|" && lastChar != "&" && lastChar != "|" {
					textQuery += " | "
				} else {
					textQuery += " "
				}
			}
			textQuery += strings.Join(strings.Fields(token), " <-> ")
		}
	}
	return
}

func extractTokens(query string) []string {
	tokens := []string{}
	quoted := strings.Split(query, "\"")
	for i, s := range quoted {
		if i%2 == 0 {
			tokens = append(tokens, strings.Fields(s)...)
		} else {
			// quoted string
			if len(tokens) > 0 {
				lastToken := tokens[len(tokens)-1]
				if len(lastToken) > 0 && lastToken[len(lastToken)-1] == ':' {
					tokens[len(tokens)-1] += s
					continue
				}
			}
			tokens = append(tokens, s)
		}
	}
	return tokens
}