github.com/mattyw/juju@v0.0.0-20140610034352-732aecd63861/state/storage/gridfs.go (about)

     1  // Copyright 2014 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package storage
     5  
     6  import (
     7  	"io"
     8  
     9  	"github.com/juju/errors"
    10  	"github.com/juju/loggo"
    11  	"labix.org/v2/mgo"
    12  )
    13  
    14  var logger = loggo.GetLogger("juju.storage")
    15  
    16  type gridFSStorage struct {
    17  	namespace string
    18  	session   *mgo.Session
    19  }
    20  
    21  var _ ResourceStorage = (*gridFSStorage)(nil)
    22  
    23  // NewGridFS returns a ResourceStorage instance backed by a mongo GridFS.
    24  // namespace is used to segregate different sets of data.
    25  func NewGridFS(namespace string, session *mgo.Session) ResourceStorage {
    26  	return &gridFSStorage{
    27  		namespace: namespace,
    28  		session:   session,
    29  	}
    30  }
    31  
    32  func (g *gridFSStorage) db() *mgo.Database {
    33  	return g.session.DB("juju")
    34  }
    35  
    36  func (g *gridFSStorage) gridFS() *mgo.GridFS {
    37  	return g.db().GridFS(g.namespace)
    38  }
    39  
    40  // Get is defined on ResourceStorage.
    41  func (g *gridFSStorage) Get(path string) (io.ReadCloser, error) {
    42  	file, err := g.gridFS().Open(path)
    43  	if err != nil {
    44  		return nil, errors.Annotatef(err, "failed to open GridFS file %q", path)
    45  	}
    46  	return file, nil
    47  }
    48  
    49  // Put is defined on ResourceStorage.
    50  func (g *gridFSStorage) Put(path string, r io.Reader, length int64) (checksum string, err error) {
    51  	file, err := g.gridFS().Create(path)
    52  	if err != nil {
    53  		return "", errors.Annotatef(err, "failed to create GridFS file %q", path)
    54  	}
    55  	defer func() {
    56  		if err != nil {
    57  			file.Close()
    58  			if removeErr := g.Remove(path); removeErr != nil {
    59  				logger.Warningf("error cleaning up after failed write: %v", removeErr)
    60  			}
    61  		}
    62  	}()
    63  	if _, err = io.CopyN(file, r, length); err != nil {
    64  		return "", errors.Annotatef(err, "failed to write data")
    65  	}
    66  	if err = file.Close(); err != nil {
    67  		return "", errors.Annotatef(err, "failed to flush data")
    68  	}
    69  	return file.MD5(), nil
    70  }
    71  
    72  // Remove is defined on ResourceStorage.
    73  func (g *gridFSStorage) Remove(path string) error {
    74  	return g.gridFS().Remove(path)
    75  }