simple-db.go

// Ivan Prisyazhnyy <[email protected]>, 2016

// Package implements simple in-memory single threaded key value database
// with nested transactions support. Interactive input support based on stdin.
// Program reads command from stdin, evaluates them and prints output to stdout.
// Errors are printed to stderr. Author missed readline emacs shortcuts while
// developing this.
package main

import (
	"bufio"
	"errors"
	"fmt"
	"os"
	"strconv"
	"strings"
)

// cmd defines command format
type cmd struct {
	op  int
	key bool
	arg bool
}

// known tokens
const (
	UNKNOWN = iota

	SET
	GET

	UNSET
	NUMEQUALTO

	BEGIN
	COMMIT
	ROLLBACK

	END
)

// cmds describes commands arguments
var cmds = map[string]cmd{
	"SET": {SET, true, true},
	"GET": {GET, true, false},

	"UNSET":      {UNSET, true, false},
	"NUMEQUALTO": {NUMEQUALTO, false, true},

	"BEGIN":    {BEGIN, false, false},
	"COMMIT":   {COMMIT, false, false},
	"ROLLBACK": {ROLLBACK, false, false},

	"END": {END, false, false},
}

// Op represents operation instance
type Op struct {
	// Token
	op int

	key string
	val int
}

// Parse a command with arguments
func Parse(cmd string) (Op, error) {
	var op Op

	s := bufio.NewScanner(strings.NewReader(cmd))
	s.Split(bufio.ScanWords)

	// Map command
	if !s.Scan() {
		return op, errors.New("command missing")
	}

	if cmd, ok := cmds[s.Text()]; !ok {
		return op, fmt.Errorf("unknown command")
	} else {
		op.op = cmd.op

		if cmd.key {
			if !s.Scan() {
				return op, errors.New("arg missing")
			}

			op.key = s.Text()
		}

		if cmd.arg {
			if !s.Scan() {
				return op, errors.New("arg missing")
			}

			var err error

			op.val, err = strconv.Atoi(s.Text())
			if err != nil {
				return op, err
			}
		}
	}

	if err := s.Err(); err != nil {
		return op, err
	}

	return op, nil
}

// Database context with binary log and transactions scopes
type DB struct {
	// Binary operations log
	binlog []Op

	// Transaction stack
	txs []int
}

func NewDB() *DB {
	return &DB{
		binlog: []Op{},
		txs: []int{},
	}
}

// Counts all keys containing requested value from the current transaction scope
func (db *DB) Count(val int) int {
	t := map[string]bool{}
	count := 0

	for i := len(db.binlog) - 1; i >= 0; i -- {
		log := db.binlog[i]
		if _, ok := t[log.key]; !ok {
			t[log.key] = true
			if log.op == SET && log.val == val {
				count ++
			}
		}
	}

	return count
}

// Lookup search for the key value current transaction observes
func (db *DB) Lookup(key string) *Op {
	for i := len(db.binlog) - 1; i >= 0; i -- {
		if db.binlog[i].key == key {
			return &db.binlog[i]
		}
	}

	return nil
}

// Executes a command on the database context
func (db *DB) Execute(op Op) error {
	switch op.op {
	case SET:
		db.binlog = append(db.binlog, op)
	case GET:
		if val := db.Lookup(op.key); val == nil || val.op != SET {
			fmt.Println("NULL")
		} else {
			fmt.Println(val.val)
		}
	case UNSET:
		db.binlog = append(db.binlog, op)
	case NUMEQUALTO:
		fmt.Println(db.Count(op.val))
	case BEGIN:
		db.txs = append(db.txs, len(db.binlog))
	case COMMIT:
		if len(db.txs) < 1 {
			return errors.New("NO TRANSACTION")
		}
		db.txs = db.txs[:len(db.txs) - 1]
	case ROLLBACK:
		if len(db.txs) < 1 {
			return errors.New("NO TRANSACTION" )
		}
		txStart := db.txs[len(db.txs) - 1]
		db.binlog = db.binlog[:txStart]
		db.txs = db.txs[:len(db.txs) - 1]
	}

	return nil
}

func main() {
	db := NewDB()

	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		cmd := strings.TrimSpace(scanner.Text())
		if len(cmd) < 1 {
			continue
		}

		op, err := Parse(cmd)
		if err != nil {
			fmt.Fprintln(os.Stderr, "invalid command:", cmd, err)
		}

		err = db.Execute(op)
		if err != nil {
			fmt.Fprintln(os.Stderr, err)
		}

		if op.op == END {
			break
		}
	}

	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
	}
}