github.com/infraboard/keyauth@v0.8.1/apps/storage/impl/storage.go (about)

     1  package impl
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  
     7  	"go.mongodb.org/mongo-driver/mongo/gridfs"
     8  	"go.mongodb.org/mongo-driver/mongo/options"
     9  
    10  	"github.com/infraboard/keyauth/apps/storage"
    11  	"github.com/infraboard/mcube/exception"
    12  )
    13  
    14  func (s *service) UploadFile(req *storage.UploadFileRequest) error {
    15  	s.log.Debugf("bucket name: %s, db file name: %s", req.BucketName, req.FileName)
    16  
    17  	if err := req.Validate(); err != nil {
    18  		return exception.NewBadRequest("valiate upload file request error, %s", err)
    19  	}
    20  
    21  	bucket, err := s.getBucket(req.BucketName)
    22  	if err != nil {
    23  		return err
    24  	}
    25  
    26  	opts := options.GridFSUpload()
    27  	opts.Metadata = req.Meta()
    28  
    29  	// 清除已有文件
    30  	bucket.Delete(req.FileName)
    31  
    32  	// 上传新文件
    33  	uploadStream, err := bucket.OpenUploadStreamWithID(req.FileName, req.FileName, opts)
    34  	if err != nil {
    35  		return err
    36  	}
    37  	defer uploadStream.Close()
    38  
    39  	fileSize, err := io.Copy(uploadStream, req.ReadCloser())
    40  	if err != nil {
    41  		return err
    42  	}
    43  
    44  	s.log.Debugf("Write file %s to DB was successful. File size: %d M", req.FileName, fileSize/1024/1024)
    45  	return nil
    46  }
    47  
    48  func (s *service) Download(req *storage.DownloadFileRequest) error {
    49  	if err := req.Validate(); err != nil {
    50  		return exception.NewBadRequest("valiate upload file request error, %s", err)
    51  	}
    52  
    53  	bucket, err := s.getBucket(req.BucketName)
    54  
    55  	s.log.Debugf("start down file: %s ...", req.FileID)
    56  	// 下载文件
    57  	size, err := bucket.DownloadToStream(req.FileID, req.Writer())
    58  	if err != nil {
    59  		return err
    60  	}
    61  
    62  	s.log.Debugf("down file: %s complete, size: %d", req.FileID, size)
    63  
    64  	return nil
    65  }
    66  
    67  func (s *service) getBucket(name string) (*gridfs.Bucket, error) {
    68  	opts := options.GridFSBucket()
    69  	opts.SetName(name)
    70  
    71  	bucket, err := gridfs.NewBucket(s.db, opts)
    72  	if err != nil {
    73  		return nil, fmt.Errorf("new bucket error, %s", err)
    74  	}
    75  
    76  	return bucket, nil
    77  }