ozu.io/shortner.go

package main

import (
	"log"
	"strconv"

	"github.com/jbenet/go-base58"
	"github.com/syndtr/goleveldb/leveldb"
)

// NewShortner returns a new Shortner.
func NewShortner() *Shortner {
	db, err := leveldb.OpenFile("shortner.db", nil)
	if err != nil {
		log.Fatal(err)
	}

	return &Shortner{
		Db:        db,
		KeyPrefix: []byte("url:"),
	}
}

// Shortner shortens URLs.
type Shortner struct {
	Db        *leveldb.DB
	KeyPrefix []byte
}

// Close calls Close() on the goleveldb.
func (s *Shortner) Close() {
	err := s.Db.Close()
	if err != nil {
		log.Fatal(err)
	}
}

// Get reads from the database.
func (s *Shortner) Get(key []byte) []byte {
	value, _ := s.Db.Get(key, nil)
	return value
}

// Set writes to the database.
func (s *Shortner) Set(key []byte, value []byte) error {
	return s.Db.Put(key, value, nil)
}

// Shorten shortens given URL
func (s *Shortner) Shorten(url []byte) ([]byte, error) {
	uid, err := s.newUID()
	if err != nil {
		return []byte{}, err
	}

	key := s.makeKey(uid)
	err = s.Set(key, url)
	if err != nil {
		return []byte{}, err
	}

	return uid, nil
}

// Lookup attempts to fetch the value for given UID
func (s *Shortner) Lookup(uid []byte) []byte {
	return s.Get(s.makeKey(uid))
}

func (s *Shortner) makeKey(uid []byte) []byte {
	return append(s.KeyPrefix, uid...)
}

func (s *Shortner) newUID() ([]byte, error) {
	key := []byte("index")

	tx, err := s.Db.OpenTransaction()
	defer tx.Commit()

	if err != nil {
		return []byte{}, err
	}

	index, _ := tx.Get(key, nil)
	if index == nil {
		index = []byte("1")
	}

	num, _ := strconv.Atoi(string(index))
	num++
	index = []byte(strconv.Itoa(num))

	err = tx.Put(key, index, nil)
	if err != nil {
		return []byte{}, err
	}

	uid := base58.Encode(index)
	return []byte(uid), nil
}