update shorturl to use badgerhold
This commit is contained in:
@@ -1,99 +1,131 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
|
||||
"github.com/coreos/bbolt"
|
||||
"github.com/gorilla/mux"
|
||||
"github.com/timshannon/badgerhold"
|
||||
|
||||
"sour.is/x/toolbox/httpsrv"
|
||||
"sour.is/x/toolbox/log"
|
||||
"sour.is/x/toolbox/uuid"
|
||||
)
|
||||
|
||||
type logger struct{}
|
||||
|
||||
func (logger) Errorf(s string, o ...interface{}) { log.Errorf(s, o...) }
|
||||
func (logger) Warningf(s string, o ...interface{}) { log.Warningf(s, o...) }
|
||||
func (logger) Infof(s string, o ...interface{}) { log.Infof(s, o...) }
|
||||
func (logger) Debugf(s string, o ...interface{}) { log.Debugf(s, o...) }
|
||||
|
||||
func init() {
|
||||
s := &shortDB{}
|
||||
httpsrv.RegisterModule("short", s.config)
|
||||
|
||||
httpsrv.HttpRegister("short", httpsrv.HttpRoutes{
|
||||
{Name: "getShort", Method: "GET", Pattern: "/s/{id}", HandlerFunc: s.getShort},
|
||||
{Name: "putShort", Method: "PUT", Pattern: "/s/{id}", HandlerFunc: s.putShort},
|
||||
{Name: "get", Method: "GET", Pattern: "/s/{id}", HandlerFunc: s.get},
|
||||
{Name: "put", Method: "PUT", Pattern: "/s/{id}", HandlerFunc: s.put},
|
||||
})
|
||||
}
|
||||
|
||||
type shortDB struct {
|
||||
path string
|
||||
bucket string
|
||||
options badgerhold.Options
|
||||
}
|
||||
|
||||
func (s *shortDB) config(config map[string]string) {
|
||||
s.bucket = "shortURL"
|
||||
if config["bucket"] != "" {
|
||||
s.bucket = config["bucket"]
|
||||
s.options = badgerhold.DefaultOptions
|
||||
s.options.Options = s.options.WithLogger(logger{})
|
||||
|
||||
var ok bool
|
||||
if s.options.Dir, ok = config["index"]; !ok {
|
||||
s.options.Dir = "data"
|
||||
}
|
||||
|
||||
s.path = "data/meta.db"
|
||||
if config["store"] != "" {
|
||||
s.path = config["store"]
|
||||
if s.options.ValueDir, ok = config["value"]; !ok {
|
||||
s.options.ValueDir = "data"
|
||||
}
|
||||
|
||||
db, err := bbolt.Open(s.path, 0666, nil)
|
||||
if err != nil {
|
||||
log.Fatalf("ShortURL: failed to open db at [%s]", s.path)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
db.Update(func(tx *bbolt.Tx) error {
|
||||
_, err := tx.CreateBucketIfNotExists([]byte(s.bucket))
|
||||
if err != nil {
|
||||
log.Fatalf("ShortURL: create bucket: %s", err)
|
||||
store, err := badgerhold.Open(s.options)
|
||||
defer func() {
|
||||
if err = store.Close(); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
return nil
|
||||
})
|
||||
|
||||
log.Noticef("ShortURL: opened db at [%s] bucket [%s]", s.path, s.bucket)
|
||||
}()
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
func (s *shortDB) getShort(w http.ResponseWriter, r *http.Request) {
|
||||
func (s *shortDB) get(w http.ResponseWriter, r *http.Request) {
|
||||
vars := mux.Vars(r)
|
||||
|
||||
id := vars["id"]
|
||||
url := s.GetURL(id)
|
||||
|
||||
if url == nil {
|
||||
httpsrv.WriteError(w, 404, "not found")
|
||||
store, err := badgerhold.Open(s.options)
|
||||
defer func() {
|
||||
if err = store.Close(); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
if err != nil {
|
||||
httpsrv.WriteText(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
var url shortURL
|
||||
err = store.Get(id, &url)
|
||||
if err != nil && err != badgerhold.ErrNotFound {
|
||||
httpsrv.WriteText(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
if err == badgerhold.ErrNotFound {
|
||||
httpsrv.WriteText(w, 404, "not found")
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Location", url.URL)
|
||||
w.WriteHeader(http.StatusFound)
|
||||
}
|
||||
|
||||
func (s *shortDB) putShort(w http.ResponseWriter, r *http.Request) {
|
||||
defer r.Body.Close()
|
||||
|
||||
func (s *shortDB) put(w http.ResponseWriter, r *http.Request) {
|
||||
vars := mux.Vars(r)
|
||||
|
||||
id := vars["id"]
|
||||
secret := r.FormValue("secret")
|
||||
u, err := url.Parse(r.FormValue("url"))
|
||||
if err != nil {
|
||||
httpsrv.WriteError(w, 400, "bad url")
|
||||
httpsrv.WriteText(w, 400, "bad url")
|
||||
return
|
||||
}
|
||||
|
||||
short := s.GetURL(id)
|
||||
store, err := badgerhold.Open(s.options)
|
||||
defer func() {
|
||||
if err = store.Close(); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
if err != nil {
|
||||
httpsrv.WriteText(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
if short == nil {
|
||||
short = newshort(id, secret, u.String())
|
||||
url := &shortURL{}
|
||||
err = store.Get(id, url)
|
||||
if err != nil && err != badgerhold.ErrNotFound {
|
||||
httpsrv.WriteText(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
s.PutURL(short.ID, short)
|
||||
httpsrv.WriteObject(w, 200, short)
|
||||
if err == badgerhold.ErrNotFound {
|
||||
url = newShortURL(id, secret, u.String())
|
||||
|
||||
err := store.Insert(url.ID, url)
|
||||
if err != nil {
|
||||
httpsrv.WriteText(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
httpsrv.WriteObject(w, 200, url)
|
||||
return
|
||||
}
|
||||
|
||||
@@ -102,24 +134,27 @@ func (s *shortDB) putShort(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
|
||||
if secret != short.Secret {
|
||||
if secret != url.Secret {
|
||||
httpsrv.WriteError(w, 403, "forbidden")
|
||||
return
|
||||
}
|
||||
|
||||
short.URL = u.String()
|
||||
err = store.Upsert(url.ID, url)
|
||||
if err != nil {
|
||||
httpsrv.WriteText(w, 500, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
s.PutURL(short.ID, short)
|
||||
httpsrv.WriteObject(w, 200, short)
|
||||
httpsrv.WriteObject(w, 200, url)
|
||||
}
|
||||
|
||||
type shortURL struct {
|
||||
ID string
|
||||
ID string `badgerhold:"unique"`
|
||||
URL string
|
||||
Secret string
|
||||
}
|
||||
|
||||
func newshort(id, secret, u string) *shortURL {
|
||||
func newShortURL(id, secret, u string) *shortURL {
|
||||
m, err := regexp.MatchString("[a-z-]{1,64}", id)
|
||||
if id == "" || !m || err != nil {
|
||||
id = uuid.V4()
|
||||
@@ -130,64 +165,3 @@ func newshort(id, secret, u string) *shortURL {
|
||||
}
|
||||
return &shortURL{ID: id, Secret: secret, URL: u}
|
||||
}
|
||||
|
||||
func (s *shortURL) Bytes() []byte {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
var w bytes.Buffer
|
||||
json.NewEncoder(&w).Encode(*s)
|
||||
return w.Bytes()
|
||||
}
|
||||
|
||||
func URLFromBytes(b []byte) *shortURL {
|
||||
if len(b) == 0 {
|
||||
return nil
|
||||
}
|
||||
var s shortURL
|
||||
json.Unmarshal(b, &s)
|
||||
|
||||
log.Debug(s)
|
||||
return &s
|
||||
}
|
||||
|
||||
func (s *shortDB) GetURL(id string) *shortURL {
|
||||
db, err := bbolt.Open(s.path, 0666, nil)
|
||||
if err != nil {
|
||||
log.Errorf("ShortURL: failed to open db at [%s]", s.path)
|
||||
return nil
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
var url *shortURL
|
||||
|
||||
err = db.View(func(tx *bbolt.Tx) error {
|
||||
b := tx.Bucket([]byte(s.bucket))
|
||||
v := b.Get([]byte(id))
|
||||
|
||||
url = URLFromBytes(v)
|
||||
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorf("ShortURL: failed to open db at [%s]", s.path)
|
||||
}
|
||||
return url
|
||||
}
|
||||
func (s *shortDB) PutURL(id string, url *shortURL) {
|
||||
db, err := bbolt.Open(s.path, 0666, nil)
|
||||
if err != nil {
|
||||
log.Errorf("ShortURL: failed to open db at [%s]", s.path)
|
||||
return
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
err = db.Update(func(tx *bbolt.Tx) error {
|
||||
b := tx.Bucket([]byte(s.bucket))
|
||||
return b.Put([]byte(id), url.Bytes())
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorf("ShortURL: failed to write db at [%s]", s.path)
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user