46 lines
843 B
Go
46 lines
843 B
Go
package shortify
|
|
|
|
import (
|
|
bolt "go.etcd.io/bbolt"
|
|
)
|
|
|
|
type Storage struct {
|
|
db *bolt.DB
|
|
bucket []byte
|
|
}
|
|
|
|
func NewStorage(db *bolt.DB) (*Storage, error) {
|
|
bucket := []byte("shorturls")
|
|
err := db.Update(func(tx *bolt.Tx) error {
|
|
_, err := tx.CreateBucketIfNotExists(bucket)
|
|
return err
|
|
})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &Storage{db: db, bucket: bucket}, nil
|
|
}
|
|
|
|
func (s *Storage) Put(id, url string) error {
|
|
return s.db.Update(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket(s.bucket)
|
|
return b.Put([]byte(id), []byte(url))
|
|
})
|
|
}
|
|
|
|
func (s *Storage) Get(id string) (string, error) {
|
|
var url []byte
|
|
err := s.db.View(func(tx *bolt.Tx) error {
|
|
b := tx.Bucket(s.bucket)
|
|
url = b.Get([]byte(id))
|
|
return nil
|
|
})
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if url == nil {
|
|
return "", ErrNotFound
|
|
}
|
|
return string(url), nil
|
|
}
|