mirror of
https://bitbucket.org/s_l_teichmann/mtsatellite
synced 2024-11-09 03:30:25 +01:00
176 lines
3.2 KiB
Go
176 lines
3.2 KiB
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/mtsatellite/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 = common.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 = common.ErrNoMoreBlocks
|
|
}
|
|
return
|
|
}
|
|
|
|
func (sbp *SQLiteBlockProducer) Close() error {
|
|
if sbp.rows != nil {
|
|
sbp.rows.Close()
|
|
}
|
|
return sbp.db.Close()
|
|
}
|