mtsatellite/interleaver/sqlite.go

// Copyright 2014 by Sascha L. Teichmann
// Use of this source code is governed by the MIT license
// that can be found in the LICENSE file.

package main

import (
	"database/sql"
	"errors"
	"os"

	"bitbucket.org/s_l_teichmann/mtredisalize/common"

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

const (
	createSql = "CREATE TABLE blocks (pos INT NOT NULL PRIMARY KEY, data BLOB)"
	insertSql = "INSERT INTO blocks (pos, data) VALUES (?, ?)"
	deleteSql = "DELETE FROM blocks"
	selectSql = "SELECT pos, data FROM blocks"
)

var ErrDatabaseNotExists = errors.New("Database does not exists.")

const blocksPerTx = 128

type (
	SQLiteBlockProducer struct {
		db       *sql.DB
		rows     *sql.Rows
		splitter common.KeySplitter
	}

	SQLiteBlockConsumer struct {
		db         *sql.DB
		insertStmt *sql.Stmt
		tx         *sql.Tx
		txCounter  int
		joiner     common.KeyJoiner
	}
)

func fileExists(path string) bool {
	_, err := os.Stat(path)
	return !os.IsNotExist(err)
}

func NewSQLiteBlockConsumer(
	path string,
	joiner common.KeyJoiner) (sbc *SQLiteBlockConsumer, err error) {

	createNew := !fileExists(path)

	var db *sql.DB
	if db, err = sql.Open("sqlite3", path); err != nil {
		return
	}

	if createNew {
		if _, err = db.Exec(createSql); err != nil {
			db.Close()
			return
		}
	} else {
		if _, err = db.Exec(deleteSql); err != nil {
			db.Close()
			return
		}
	}

	var insertStmt *sql.Stmt
	if insertStmt, err = db.Prepare(insertSql); err != nil {
		db.Close()
		return
	}

	var tx *sql.Tx
	if tx, err = db.Begin(); err != nil {
		insertStmt.Close()
		db.Close()
		return
	}

	sbc = &SQLiteBlockConsumer{
		db:         db,
		insertStmt: insertStmt,
		tx:         tx,
		joiner:     joiner}

	return
}

func (sbc *SQLiteBlockConsumer) Close() error {
	sbc.tx.Commit()
	sbc.insertStmt.Close()
	return sbc.db.Close()
}

func (sbc *SQLiteBlockConsumer) getTx() (tx *sql.Tx, err error) {
	if sbc.txCounter >= blocksPerTx {
		sbc.txCounter = 0
		if err = sbc.tx.Commit(); err != nil {
			return
		}
		if sbc.tx, err = sbc.db.Begin(); err != nil {
			return
		}
	}
	sbc.txCounter++
	tx = sbc.tx
	return
}

func (sbc *SQLiteBlockConsumer) Consume(block *common.Block) (err error) {
	var tx *sql.Tx
	if tx, err = sbc.getTx(); err != nil {
		return
	}
	_, err = tx.Stmt(sbc.insertStmt).Exec(sbc.joiner(block.Coord), block.Data)
	return
}

func NewSQLiteBlockProducer(
	path string,
	splitter common.KeySplitter) (sbp *SQLiteBlockProducer, err error) {

	if !fileExists(path) {
		err = ErrDatabaseNotExists
		return
	}

	var db *sql.DB
	if db, err = sql.Open("sqlite3", path); err != nil {
		return
	}

	var rows *sql.Rows
	if rows, err = db.Query(selectSql); err != nil {
		db.Close()
		return
	}

	sbp = &SQLiteBlockProducer{
		db:       db,
		rows:     rows,
		splitter: splitter}

	return
}

func (sbp *SQLiteBlockProducer) Next(block *common.Block) (err error) {
	if sbp.rows == nil {
		err = ErrNoMoreBlocks
		return
	}
	if sbp.rows.Next() {
		var key int64
		if err = sbp.rows.Scan(&key, &block.Data); err == nil {
			block.Coord = sbp.splitter(key)
		}
	} else {
		sbp.rows.Close()
		sbp.rows = nil
		err = ErrNoMoreBlocks
	}
	return
}

func (sbp *SQLiteBlockProducer) Close() error {
	if sbp.rows != nil {
		sbp.rows.Close()
	}
	return sbp.db.Close()
}