qurl/storage/bolt/qurl.go

package bolt

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

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

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

func (stor *BoltStorage) AddQURL(qurl *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 qurl.ID == 0 {
		qurl.ID, err = rb.NextSequence()
		if err != nil {
			return err
		}
	}

	// Create a byte array from the ID
	bid := make([]byte, 8)
	binary.BigEndian.PutUint64(bid, qurl.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(qurl.URL), bid)

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

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

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

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

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

	if len(qurl.Browser) > 0 {
		err = qb.Put(BrowserField, []byte(qurl.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
	}

	qb.SetSequence(seq)

	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
	}

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

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

	return qurl, 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()

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

	bu := []byte(u)
	rc := rb.Cursor()
	for k, _ := rc.First(); k != nil; k, _ = rc.Next() {
		qb := rb.Bucket(k)
		if qb == nil {
			continue
		}

		qu := qb.Get(URLField)
		if bytes.Equal(bu, qu) {
			qurl := &qurl.QURL{
				ID: binary.BigEndian.Uint64(k),
				URL: string(qu),
			}
			return qurl, nil
		}
	}
	return nil, nil
}
*/