yatwiki/DB.go

package yatwiki

import (
	"database/sql"
	"fmt"
	"log"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type WikiDB struct {
	db               *sql.DB
	compressionLevel int
}

func NewWikiDB(dbFilePath string, compressionLevel int) (*WikiDB, error) {
	db, err := sql.Open("sqlite3", dbFilePath)
	if err != nil {
		return nil, err
	}

	wdb := WikiDB{
		db:               db,
		compressionLevel: compressionLevel,
	}

	err = wdb.assertSchema()
	if err != nil {
		return nil, fmt.Errorf("assertSchema: %s", err.Error())
	}

	return &wdb, nil
}

func (this *WikiDB) multiTx(stmts ...string) (err error) {
	tx, err := this.db.Begin()
	if err != nil {
		return err
	}

	for _, stmt := range stmts {
		_, err := tx.Exec(stmt)
		if err != nil {
			tx.Rollback()
			return err
		}
	}

	return tx.Commit()
}

func (this *WikiDB) assertSchema() error {
	_, err := this.db.Exec(`
	CREATE TABLE IF NOT EXISTS schema (
		id INTEGER PRIMARY KEY
	);`)
	if err != nil {
		return err
	}

	// Look up current value from schema table
	schemaLookup := this.db.QueryRow(`SELECT MAX(id) mid FROM schema`)
	currentSchema := int64(0)
	if err = schemaLookup.Scan(&currentSchema); err == nil {
		// That's fine
	}

	log.Printf("Found DB version %d\n", currentSchema)

	//

	if currentSchema == 0 {
		// Schema 0 ==> Schema 1
		log.Println("Upgrading to DB version 1")

		err := this.multiTx(
			`CREATE TABLE IF NOT EXISTS articles (
				id			INTEGER PRIMARY KEY,
				article		INTEGER,
				modified	INTEGER,
				body		BLOB,
				author		TEXT
			);`,
			`CREATE TABLE IF NOT EXISTS titles (
				id INTEGER PRIMARY KEY,
				title TEXT
			);`,
			`CREATE INDEX IF NOT EXISTS articles_modified_index ON articles (modified)`,
			`CREATE INDEX IF NOT EXISTS articles_title_index ON articles (article)`,
			`INSERT INTO schema (id) VALUES (1);`,
		)
		if err != nil {
			return err
		}
		currentSchema = 1
	}

	//

	if currentSchema == 1 {
		log.Println("Upgrading to DB version 2")

		err := this.multiTx(
			`ALTER TABLE titles ADD COLUMN is_deleted INTEGER DEFAULT 0;`,
			`INSERT INTO schema (id) VALUES (2);`,
		)
		if err != nil {
			return err
		}
		currentSchema = 1
	}

	//

	return nil
}

type TitleInfo struct {
	TitleID   int64
	Title     string
	IsDeleted bool
}

type Article struct {
	TitleInfo
	ArticleID int64
	Modified  int64
	Body      []byte
	Author    string
}

func (this *WikiDB) GetArticleById(articleId int) (*Article, error) {
	row := this.db.QueryRow(`SELECT articles.* FROM articles WHERE id = ?`, articleId)
	return this.parseArticle(row)
}

func (this *WikiDB) GetRevision(revId int) (*Article, error) {
	row := this.db.QueryRow(`SELECT articles.*, titles.title, titles.is_deleted FROM articles JOIN titles ON articles.article=titles.id WHERE articles.id = ?`, revId)
	return this.parseArticleWithTitle(row)
}

func (this *WikiDB) GetLatestVersion(title string) (*Article, error) {
	row := this.db.QueryRow(`SELECT articles.* FROM articles WHERE article = (SELECT id FROM titles WHERE title = ?) ORDER BY modified DESC LIMIT 1`, this.normaliseTitle(title))
	return this.parseArticle(row)
}

type ArticleAlteredError struct {
	got, expected int64
}

func (aae ArticleAlteredError) Error() string {
	return fmt.Sprintf("Warning: Your changes were not based on the most recent version of the page (r%d ≠ r%d). No changes were saved.", aae.got, aae.expected)
}

func (this *WikiDB) normaliseTitle(title string) string {
	return strings.ToLower(strings.Trim(title, " \r\n\t"))
}

func (this *WikiDB) SaveArticle(title, author, body string, expectBaseRev int64) error {
	isNewArticle := false
	a, err := this.GetLatestVersion(title)
	if err != nil {
		if err == sql.ErrNoRows {
			isNewArticle = true
		} else {
			return fmt.Errorf("Couldn't check for existing article title: %s", err.Error())
		}
	}

	if !isNewArticle && a.ArticleID != expectBaseRev {
		return ArticleAlteredError{got: expectBaseRev, expected: a.ArticleID}
	}

	zBody, err := gzdeflate([]byte(body), this.compressionLevel)
	if err != nil {
		return err
	}

	var titleId int64
	if isNewArticle {
		titleInsert, err := this.db.Exec(`INSERT INTO titles (title) VALUES (?)`, this.normaliseTitle(title))
		if err != nil {
			return err
		}

		titleId, err = titleInsert.LastInsertId()
		if err != nil {
			return err
		}
	} else {
		titleId = a.TitleID
	}

	_, err = this.db.Exec(`INSERT INTO articles (article, modified, body, author) VALUES (?, ?, ?, ?)`, titleId, time.Now().Unix(), zBody, author)
	if err != nil {
		return err
	}

	// Update is-deleted flag
	isDeleted := 0
	if body == bbcodeIsDeleted {
		isDeleted = 1
	}
	_, err = this.db.Exec(`UPDATE titles SET is_deleted = ? WHERE id = ?`, isDeleted, titleId)

	return nil
}

func (this *WikiDB) GetRevisionHistory(title string) ([]Article, error) {
	rows, err := this.db.Query(
		`SELECT articles.id, articles.modified, articles.author FROM articles WHERE article = (SELECT id FROM titles WHERE title = ?) ORDER BY modified DESC`,
		this.normaliseTitle(title),
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ret := make([]Article, 0)
	for rows.Next() {
		a := Article{}
		err := rows.Scan(&a.ArticleID, &a.Modified, &a.Author)
		if err != nil {
			return nil, err
		}
		ret = append(ret, a)
	}
	return ret, nil
}

func (this *WikiDB) GetNextOldestRevision(revision int) (int, error) {
	row := this.db.QueryRow(
		`SELECT articles.id FROM articles WHERE articles.article = (SELECT article FROM articles WHERE id = ?) AND id < ? ORDER BY id DESC LIMIT 1`,
		revision, revision,
	)

	var ret int
	err := row.Scan(&ret)
	if err != nil {
		return 0, err
	}

	return ret, nil
}

func (this *WikiDB) GetRecentChanges(offset int, limit int) ([]Article, error) {
	rows, err := this.db.Query(
		`SELECT articles.id, articles.modified, articles.author, titles.title, titles.is_deleted FROM articles JOIN titles ON articles.article=titles.id ORDER BY modified DESC ` +
			fmt.Sprintf(`LIMIT %d OFFSET %d`, limit, offset),
	)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ret := make([]Article, 0, limit)
	for rows.Next() {
		a := Article{}
		err := rows.Scan(&a.ArticleID, &a.Modified, &a.Author, &a.Title, &a.IsDeleted)
		if err != nil {
			return nil, err
		}
		ret = append(ret, a)
	}
	return ret, nil
}

func (this *WikiDB) TotalRevisions() (int64, error) {
	row := this.db.QueryRow(`SELECT COUNT(*) c FROM articles`)
	var ret int64
	err := row.Scan(&ret)
	if err != nil {
		return 0, err
	}

	return ret, nil
}

func (this *WikiDB) ListTitles(includeDeleted bool) ([]TitleInfo, error) {
	stmt := `SELECT id, title, is_deleted FROM titles`
	if !includeDeleted {
		stmt += ` WHERE is_deleted = 0`
	}
	stmt += ` ORDER BY title ASC`

	rows, err := this.db.Query(stmt)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ret := make([]TitleInfo, 0)
	for rows.Next() {
		var title TitleInfo
		err = rows.Scan(&title.TitleID, &title.Title, &title.IsDeleted)
		if err != nil {
			return nil, err
		}

		ret = append(ret, title)
	}
	return ret, nil
}

func (this *WikiDB) parseArticle(row *sql.Row) (*Article, error) {
	a := Article{}
	var gzBody []byte
	err := row.Scan(&a.ArticleID, &a.TitleID, &a.Modified, &gzBody, &a.Author)
	if err != nil {
		return nil, err
	}

	decompressed, err := gzinflate(gzBody)
	if err != nil {
		return nil, err
	}

	a.Body = decompressed

	return &a, nil
}

func (this *WikiDB) parseArticleWithTitle(row *sql.Row) (*Article, error) {
	a := Article{}
	var gzBody []byte
	err := row.Scan(&a.ArticleID, &a.TitleID, &a.Modified, &gzBody, &a.Author, &a.Title, &a.IsDeleted)
	if err != nil {
		return nil, err
	}

	decompressed, err := gzinflate(gzBody)
	if err != nil {
		return nil, err
	}

	a.Body = decompressed

	return &a, nil
}

func (this *WikiDB) Close() {
	this.db.Close()
}