|
| 1 | +package baidubos |
| 2 | + |
| 3 | +import ( |
| 4 | + "fmt" |
| 5 | + "io" |
| 6 | + "net/http" |
| 7 | + "strings" |
| 8 | + |
| 9 | + "github.com/baidu/baiducloud-sdk-go/bce" |
| 10 | + "github.com/baidu/baiducloud-sdk-go/bos" |
| 11 | + oss "github.com/langgenius/dify-cloud-kit/oss" |
| 12 | +) |
| 13 | + |
| 14 | +const OSS_TYPE_BAIDU_BOS = "baidu_bos" |
| 15 | + |
| 16 | +type bosClient interface { |
| 17 | + PutObject(bucketName string, objectKey string, data interface{}, metadata *bos.ObjectMetadata, option *bce.SignOption) (bos.PutObjectResponse, error) |
| 18 | + GetObject(bucketName string, objectKey string, option *bce.SignOption) (*bos.Object, error) |
| 19 | + GetObjectMetadata(bucketName string, objectKey string, option *bce.SignOption) (*bos.ObjectMetadata, error) |
| 20 | + ListObjectsFromRequest(req bos.ListObjectsRequest, option *bce.SignOption) (*bos.ListObjectsResponse, error) |
| 21 | + DeleteObject(bucketName string, objectKey string, option *bce.SignOption) error |
| 22 | +} |
| 23 | + |
| 24 | +type BaiduBOSStorage struct { |
| 25 | + bucket string |
| 26 | + client bosClient |
| 27 | +} |
| 28 | + |
| 29 | +type BaiduBOSConfig struct { |
| 30 | + Endpoint string |
| 31 | + AccessKey string |
| 32 | + SecretKey string |
| 33 | + Region string |
| 34 | + Bucket string |
| 35 | +} |
| 36 | + |
| 37 | +func (c *BaiduBOSConfig) Validate() error { |
| 38 | + if c.Bucket == "" || c.AccessKey == "" || c.SecretKey == "" { |
| 39 | + return fmt.Errorf("bucket, accessKey, secretKey cannot be empty") |
| 40 | + } |
| 41 | + if c.Endpoint == "" && c.Region == "" { |
| 42 | + return fmt.Errorf("endpoint or region must be provided") |
| 43 | + } |
| 44 | + return nil |
| 45 | +} |
| 46 | + |
| 47 | +func NewBaiduBOSStorage(config BaiduBOSConfig) (oss.OSS, error) { |
| 48 | + config.Endpoint = strings.TrimSpace(config.Endpoint) |
| 49 | + config.Region = strings.TrimSpace(config.Region) |
| 50 | + config.AccessKey = strings.TrimSpace(config.AccessKey) |
| 51 | + config.SecretKey = strings.TrimSpace(config.SecretKey) |
| 52 | + config.Bucket = strings.TrimSpace(config.Bucket) |
| 53 | + |
| 54 | + if err := config.Validate(); err != nil { |
| 55 | + return nil, oss.ErrArgumentInvalid.WithDetail(err.Error()) |
| 56 | + } |
| 57 | + |
| 58 | + credentials := bce.NewCredentials(config.AccessKey, config.SecretKey) |
| 59 | + bceConfig := bce.NewConfig(credentials) |
| 60 | + |
| 61 | + if config.Endpoint != "" { |
| 62 | + bceConfig.Endpoint = config.Endpoint |
| 63 | + } |
| 64 | + if config.Region != "" { |
| 65 | + bceConfig.Region = config.Region |
| 66 | + } |
| 67 | + |
| 68 | + client := bos.NewClient(bceConfig) |
| 69 | + |
| 70 | + return newBaiduBOSStorageWithClient(config.Bucket, client), nil |
| 71 | +} |
| 72 | + |
| 73 | +func newBaiduBOSStorageWithClient(bucket string, client bosClient) *BaiduBOSStorage { |
| 74 | + return &BaiduBOSStorage{ |
| 75 | + bucket: bucket, |
| 76 | + client: client, |
| 77 | + } |
| 78 | +} |
| 79 | + |
| 80 | +func (b *BaiduBOSStorage) Save(key string, data []byte) error { |
| 81 | + _, err := b.client.PutObject(b.bucket, key, data, nil, nil) |
| 82 | + return err |
| 83 | +} |
| 84 | + |
| 85 | +func (b *BaiduBOSStorage) Load(key string) ([]byte, error) { |
| 86 | + obj, err := b.client.GetObject(b.bucket, key, nil) |
| 87 | + if err != nil { |
| 88 | + return nil, err |
| 89 | + } |
| 90 | + defer obj.ObjectContent.Close() |
| 91 | + |
| 92 | + return io.ReadAll(obj.ObjectContent) |
| 93 | +} |
| 94 | + |
| 95 | +func (b *BaiduBOSStorage) Exists(key string) (bool, error) { |
| 96 | + _, err := b.client.GetObjectMetadata(b.bucket, key, nil) |
| 97 | + if err != nil { |
| 98 | + if isNotFoundError(err) { |
| 99 | + return false, nil |
| 100 | + } |
| 101 | + return false, err |
| 102 | + } |
| 103 | + return true, nil |
| 104 | +} |
| 105 | + |
| 106 | +func (b *BaiduBOSStorage) State(key string) (oss.OSSState, error) { |
| 107 | + metadata, err := b.client.GetObjectMetadata(b.bucket, key, nil) |
| 108 | + if err != nil { |
| 109 | + return oss.OSSState{}, err |
| 110 | + } |
| 111 | + |
| 112 | + return oss.OSSState{ |
| 113 | + Size: metadata.ContentLength, |
| 114 | + }, nil |
| 115 | +} |
| 116 | + |
| 117 | +func (b *BaiduBOSStorage) List(prefix string) ([]oss.OSSPath, error) { |
| 118 | + if prefix != "" && !strings.HasSuffix(prefix, "/") { |
| 119 | + prefix = prefix + "/" |
| 120 | + } |
| 121 | + |
| 122 | + marker := "" |
| 123 | + paths := []oss.OSSPath{} |
| 124 | + for { |
| 125 | + req := bos.ListObjectsRequest{ |
| 126 | + BucketName: b.bucket, |
| 127 | + Prefix: prefix, |
| 128 | + Marker: marker, |
| 129 | + MaxKeys: 1000, |
| 130 | + } |
| 131 | + resp, err := b.client.ListObjectsFromRequest(req, nil) |
| 132 | + if err != nil { |
| 133 | + return nil, err |
| 134 | + } |
| 135 | + |
| 136 | + for _, obj := range resp.Contents { |
| 137 | + key := strings.TrimPrefix(obj.Key, prefix) |
| 138 | + key = strings.TrimPrefix(key, "/") |
| 139 | + |
| 140 | + if key == "" { |
| 141 | + continue |
| 142 | + } |
| 143 | + paths = append(paths, oss.OSSPath{ |
| 144 | + Path: key, |
| 145 | + IsDir: false, |
| 146 | + }) |
| 147 | + } |
| 148 | + |
| 149 | + if !resp.IsTruncated { |
| 150 | + break |
| 151 | + } |
| 152 | + |
| 153 | + if resp.NextMarker != "" { |
| 154 | + marker = resp.NextMarker |
| 155 | + } else if len(resp.Contents) > 0 { |
| 156 | + marker = resp.Contents[len(resp.Contents)-1].Key |
| 157 | + } else { |
| 158 | + break |
| 159 | + } |
| 160 | + } |
| 161 | + |
| 162 | + return paths, nil |
| 163 | +} |
| 164 | + |
| 165 | +func (b *BaiduBOSStorage) Delete(key string) error { |
| 166 | + return b.client.DeleteObject(b.bucket, key, nil) |
| 167 | +} |
| 168 | + |
| 169 | +func (b *BaiduBOSStorage) Type() string { |
| 170 | + return OSS_TYPE_BAIDU_BOS |
| 171 | +} |
| 172 | + |
| 173 | +func isNotFoundError(err error) bool { |
| 174 | + if bceErr, ok := err.(*bce.Error); ok { |
| 175 | + return bceErr.StatusCode == http.StatusNotFound |
| 176 | + } |
| 177 | + return false |
| 178 | +} |
0 commit comments