qurl/storage/bolt/qurl.go

package bolt

import (
	"encoding/binary"
	"qurl/qurl"
)

var (
	QURLBucket    = []byte{0x00}
	ReverseBucket = []byte{0x01}

	URLField     = []byte{0x00}
	CreatedField = []byte{0x01}
	IPField      = []byte{0x02}
	BrowserField = []byte{0x03}
)

func (stor *BoltStorage) AddQURL(q *qurl.QURL) error {
	tx, err := stor.DB.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	rb, err := tx.CreateBucketIfNotExists(QURLBucket)
	if err != nil {
		return err
	}

	// Populate the ID from the sequence if we don't have one
	if q.ID == 0 {
		s, err := rb.NextSequence()
		if err != nil {
			return err
		}

		// Bolt's sequence starts at 1, so for
		// backwards compatibility we have subtract
		// one so we're zero-based
		q.ID = (s - 1)
	}

	// Create a byte array from the ID
	bid := make([]byte, 8)
	binary.BigEndian.PutUint64(bid, q.ID)

	// Add an entry into the reverse indexed bucket for quickly
	// determining if a URL is already in the database
	ab, err := tx.CreateBucketIfNotExists(ReverseBucket)
	if err != nil {
		return err
	}
	err = ab.Put([]byte(q.URL), bid)

	qb, err := rb.CreateBucketIfNotExists(bid)
	if err != nil {
		return err
	}

	// Write the ID to URL
	err = qb.Put(URLField, []byte(q.URL))
	if err != nil {
		return err
	}

	if !q.Created.IsZero() {
		// Create byte array from the Created date
		bdt, err := q.Created.MarshalBinary()
		if err != nil {
			return err
		}

		// Write the Created date
		err = qb.Put(CreatedField, bdt)
		if err != nil {
			return err
		}
	}

	if q.IP != nil {
		err = qb.Put(IPField, q.IP)
		if err != nil {
			return err
		}
	}

	if len(q.Browser) > 0 {
		err = qb.Put(BrowserField, []byte(q.Browser))
		if err != nil {
			return err
		}
	}

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

	return nil
}

func (stor *BoltStorage) SetQURLSequence(seq uint64) error {
	tx, err := stor.DB.Begin(true)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	qb, err := tx.CreateBucketIfNotExists(QURLBucket)
	if err != nil {
		return err
	}

	// Since the sequence number is decremented by one
	// for backwards compatibility (see above)
	// we increment it by one when setting the sequence
	qb.SetSequence(seq + 1)

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

	return nil
}

func (stor *BoltStorage) GetQURLByURL(u string) (*qurl.QURL, error) {
	tx, err := stor.DB.Begin(false)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	ab := tx.Bucket(ReverseBucket)
	if ab == nil {
		return nil, nil
	}

	bid := ab.Get([]byte(u))
	if bid == nil {
		return nil, nil
	}

	rb := tx.Bucket(QURLBucket)
	if rb == nil {
		return nil, nil
	}

	qb := rb.Bucket(bid)
	if qb == nil {
		return nil, nil
	}

	q := &qurl.QURL{ID: binary.BigEndian.Uint64(bid)}

	qu := qb.Get(URLField)
	if qu != nil {
		q.URL = string(qu)
	}

	return q, nil
}

func (stor *BoltStorage) GetQURLByID(id uint64) (*qurl.QURL, error) {
	tx, err := stor.DB.Begin(false)
	if err != nil {
		return nil, err
	}
	defer tx.Rollback()

	// Create a byte array from the ID
	bid := make([]byte, 8)
	binary.BigEndian.PutUint64(bid, id)

	rb := tx.Bucket(QURLBucket)
	if rb == nil {
		return nil, nil
	}

	qb := rb.Bucket(bid)
	if qb == nil {
		return nil, nil
	}

	q := &qurl.QURL{ID: id}

	qu := qb.Get(URLField)
	if qu != nil {
		q.URL = string(qu)
	}

	return q, nil
}