Initial commit

storage/bolt/bolt.go

@@ -0,0 +1,28 @@
+package bolt
+
+import (
+	bolt "go.etcd.io/bbolt"
+	"net/url"
+	"time"
+)
+
+type BoltStorage struct {
+	DB *bolt.DB
+}
+
+func (stor *BoltStorage) Shutdown() {
+	stor.DB.Close()
+}
+
+func New(u *url.URL) (*BoltStorage, error) {
+	path := u.Opaque
+	if u.Path != "" {
+		path = u.Path
+	}
+
+	db, err := bolt.Open(path, 0600, &bolt.Options{Timeout: 3 * time.Second})
+	if err != nil {
+		return nil, err
+	}
+	return &BoltStorage{DB: db}, nil
+}

storage/bolt/qurl.go

@@ -0,0 +1,188 @@
+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
+}
+*/