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 }