gosh/db.go

package main

import (
	"database/sql"
	"fmt"
	"os"

	_ "code.google.com/p/gosqlite/sqlite3" // registers sqlite
)

type Database struct {
	conn *sql.DB
}

func (db *Database) Close() error {
	return db.conn.Close()
}

func (db Database) GetCommand(commandStr string) (*Command, error) {
	rows, err := db.conn.Query("SELECT rowid, command FROM commands WHERE command = ?;", commandStr)
	if err != nil {
		return nil, err
	}
	var (
		rowid   int64
		command string
	)
	for rows.Next() {
		err = rows.Scan(&rowid, &command)
		if err != nil {
			rows.Close()
			return nil, err
		}
	}
	rows.Close()
	if rowid == 0 {
		return nil, nil
	}
	cmd := Command{rowid: rowid, Args: command}

	if rows, err = db.conn.Query("SELECT time FROM history WHERE command_id = ?;", rowid); err != nil {
		return nil, err
	}
	for rows.Next() {
		var t int64
		if err = rows.Scan(&t); err != nil {
			rows.Close()
			return nil, err
		}
		cmd.Times = append(cmd.Times, t)
	}
	rows.Close()

	return &cmd, nil
}

func (db *Database) InsertCommand(cmdTime int64, command string) error {
	var (
		cmd *Command
		err error
	)

	// will be
	if cmd, err = db.GetCommand(command); err != nil {
		return err
	}
	fmt.Printf("%#v\n", cmd)

	insertCommand, err := db.conn.Prepare("INSERT INTO commands (command) VALUES (?);")
	if err != nil {
		return err
	}
	insertTime, err := db.conn.Prepare("INSERT INTO history (command_id, time) VALUES (?,?);")
	if err != nil {
		return err
	}

	tx, err := db.conn.Begin()
	if err != nil {
		return err
	}

	var (
		commandId int64
		hasTime   bool
	)
	if cmd == nil {
		res, err := tx.Stmt(insertCommand).Exec(command)
		if err != nil {
			tx.Rollback()
			return err
		}
		commandId, err = res.LastInsertId()
		if err != nil {
			tx.Rollback()
			return err
		}
		cmd = &Command{rowid: commandId, Args: command}
	} else {
		commandId = cmd.rowid
	}
	for i := range cmd.Times {
		if cmd.Times[i] == cmdTime {
			hasTime = true
		}
	}
	if !hasTime {
		if _, err = tx.Stmt(insertTime).Exec(commandId, cmdTime); err != nil {
			tx.Rollback()
			return err
		}
	}

	if err := tx.Commit(); err != nil {
		return err
	}

	return nil
}

func NewDatabase(path string) (*Database, error) {
	initDb := false

	stat, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			initDb = true
		} else {
			return nil, err
		}
	}
	if stat != nil && stat.Size() == 0 {
		initDb = true
	}

	conn, err := sql.Open("sqlite3", path)
	if err != nil {
		return nil, err
	}
	db := &Database{conn: conn}

	if initDb {
		if err := initializeDatabase(db); err != nil {
			return nil, err
		}
	}

	return db, nil
}

func initializeDatabase(db *Database) error {
	if _, err := db.conn.Exec(createCommandsTable); err != nil {
		return err
	}
	if _, err := db.conn.Exec(createHistoryTable); err != nil {
		return err
	}
	if _, err := db.conn.Exec(createIndexes); err != nil {
		return err
	}

	return nil
}

const (
	//id int NOT NULL PRIMARY KEY,
	createCommandsTable = `
    CREATE TABLE IF NOT EXISTS commands (
      command text NOT NULL
    ); `
	//id text NOT NULL PRIMARY KEY,
	createHistoryTable = `
    CREATE TABLE IF NOT EXISTS history (
      "command_id" int NOT NULL,
      "time" bigint NOT NULL,
      CONSTRAINT "command_fk" FOREIGN KEY ("command_id") REFERENCES "commands" ("id")
    ); `
	createIndexes = `
    CREATE UNIQUE INDEX IF NOT EXISTS "commands_ix" ON "commands" (command);
  `
)