yatwiki/DB.go

package yatwiki3

import (
	"bytes"
	"compress/flate"
	"crypto/md5"
	"database/sql"
	"encoding/hex"
	"fmt"
	"io"
	"net/http"
	"time"

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

type WikiDB struct {
	db *sql.DB
}

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

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

	return &wdb, nil
}

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

	_, err = this.db.Exec(`
	CREATE TABLE IF NOT EXISTS titles (
		id INTEGER PRIMARY KEY,
		title TEXT
	);`)
	if err != nil {
		return err
	}

	_, err = this.db.Exec(`
	CREATE INDEX IF NOT EXISTS articles_modified_index ON articles (modified)
	`)
	if err != nil {
		return err
	}

	_, err = this.db.Exec(`
	CREATE INDEX IF NOT EXISTS articles_title_index ON articles (article)
	`)
	if err != nil {
		return err
	}

	return nil
}

type Article struct {
	ID       int
	TitleID  int
	Modified int64
	Body     []byte
	Author   string
}

func (this *Article) FillModifiedTimestamp() {
	this.Modified = time.Now().Unix()
}

func (this *Article) FillAuthor(r *http.Request) {
	userAgentHash := md5.Sum([]byte(r.UserAgent()))
	this.Author = r.RemoteAddr + "-" + hex.EncodeToString(userAgentHash[:])[:6]
}

type ArticleWithTitle struct {
	Article
	Title 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) 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`, title)
	return this.parseArticle(row)
}

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() ([]string, error) {
	rows, err := this.db.Query(`SELECT title FROM titles ORDER BY title ASC`)
	if err != nil {
		return nil, err
	}

	defer rows.Close()
	ret := make([]string, 0)
	for rows.Next() {
		var title string
		err = rows.Scan(&title)
		if err != nil {
			return nil, err
		}

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

func (this *WikiDB) gzinflate(gzBody []byte) ([]byte, error) {
	gzBodyReader := bytes.NewReader(gzBody)

	gzReader := flate.NewReader(gzBodyReader)
	defer gzReader.Close()

	buffer := bytes.Buffer{}
	_, err := io.Copy(&buffer, gzReader)
	if err != nil {
		return nil, err
	}

	return buffer.Bytes(), nil
}

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

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

	a.Body = decompressed

	return &a, nil
}

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

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

	a.Body = decompressed

	return &a, nil
}

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