golang-url-shortener/store/store.go

// Package store provides support to intergate with the entries
package store

import (
	"encoding/json"
	"math/rand"
	"time"

	"github.com/asaskevich/govalidator"
	"github.com/boltdb/bolt"
	"github.com/pkg/errors"
)

// Store holds internal funcs and vars about the store
type Store struct {
	db         *bolt.DB
	bucketName []byte
	idLength   uint
}

// Entry is the data set which is stored in the DB as JSON
type Entry struct {
	URL        string
	VisitCount int
	RemoteAddr string `json:",omitempty"`
	CreatedOn  time.Time
	LastVisit  time.Time
}

// ErrNoEntryFound is returned when no entry to a id is found
var ErrNoEntryFound = errors.New("no entry found")

// ErrNoValidURL is returned when the URL is not valid
var ErrNoValidURL = errors.New("no valid URL")

// ErrGeneratingTriesFailed is returned when the 10 tries to generate an id failed
var ErrGeneratingTriesFailed = errors.New("could not generate unique id")

// ErrIDIsEmpty is returned when the given ID is empty
var ErrIDIsEmpty = errors.New("id is empty")

// New initializes the store with the db
func New(dbName string, idLength uint) (*Store, error) {
	db, err := bolt.Open(dbName, 0644, &bolt.Options{Timeout: 1 * time.Second})
	if err != nil {
		return nil, errors.Wrap(err, "could not open bolt DB database")
	}
	bucketName := []byte("shortUrlBkt")
	err = db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists(bucketName)
		return err
	})
	if err != nil {
		return nil, err
	}
	return &Store{
		db:         db,
		idLength:   idLength,
		bucketName: bucketName,
	}, nil
}

// GetEntryByID returns a unmarshalled entry of the db by a given ID
func (s *Store) GetEntryByID(id string) (*Entry, error) {
	if id == "" {
		return nil, ErrIDIsEmpty
	}
	raw, err := s.GetEntryByIDRaw(id)
	if err != nil {
		return nil, err
	}
	var entry *Entry
	err = json.Unmarshal(raw, &entry)
	return entry, err
}

// IncreaseVisitCounter increments the visit counter of an entry
func (s *Store) IncreaseVisitCounter(id string) error {
	entry, err := s.GetEntryByID(id)
	if err != nil {
		return err
	}
	entry.VisitCount++
	entry.LastVisit = time.Now()
	raw, err := json.Marshal(entry)
	if err != nil {
		return err
	}
	err = s.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(s.bucketName)
		err := bucket.Put([]byte(id), raw)
		if err != nil {
			return errors.Wrap(err, "could not put data into bucket")
		}
		return nil
	})
	return err
}

// GetEntryByIDRaw returns the raw data (JSON) of a data set
func (s *Store) GetEntryByIDRaw(id string) ([]byte, error) {
	var raw []byte
	err := s.db.View(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(s.bucketName)
		raw = bucket.Get([]byte(id))
		if raw == nil {
			return ErrNoEntryFound
		}
		return nil
	})
	return raw, err
}

// CreateEntry creates a new record and returns his short id
func (s *Store) CreateEntry(URL, remoteAddr string) (string, error) {
	if !govalidator.IsURL(URL) {
		return "", ErrNoValidURL
	}
	// try it 10 times to make a short URL
	for i := 1; i <= 10; i++ {
		id, err := s.createEntry(URL, remoteAddr)
		if err != nil {
			continue
		}
		return id, nil
	}
	return "", ErrGeneratingTriesFailed
}

// checkExistens returns true if a entry with a given ID
// exists and false if not
func (s *Store) checkExistence(id string) bool {
	raw, err := s.GetEntryByIDRaw(id)
	if err != nil && err != ErrNoEntryFound {
		return true
	}
	if raw != nil {
		return true
	}
	return false
}

// createEntry creates a new entry
func (s *Store) createEntry(URL, remoteAddr string) (string, error) {
	id := generateRandomString(s.idLength)

	exists := s.checkExistence(id)
	if !exists {
		raw, err := json.Marshal(Entry{
			URL:        URL,
			RemoteAddr: remoteAddr,
			CreatedOn:  time.Now(),
		})
		if err != nil {
			return "", err
		}
		return id, s.createEntryRaw([]byte(id), raw)
	}
	return "", ErrGeneratingTriesFailed
}

// createEntryRaw creates a entry with the given key value pair
func (s *Store) createEntryRaw(key, value []byte) error {
	err := s.db.Update(func(tx *bolt.Tx) error {
		bucket := tx.Bucket(s.bucketName)
		raw := bucket.Get(key)
		if raw != nil {
			return errors.New("entry value is not empty")
		}
		err := bucket.Put(key, value)
		if err != nil {
			return errors.Wrap(err, "could not put data into bucket")
		}
		return nil
	})
	return err
}

// Close closes the bolt db database
func (s *Store) Close() error {
	return s.db.Close()
}

// generateRandomString generates a random string with an predefined length
func generateRandomString(length uint) string {
	letterRunes := []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
	b := make([]rune, length)
	for i := range b {
		b[i] = letterRunes[rand.Intn(len(letterRunes))]
	}
	return string(b)
}