ossService/oss/oss.go

package oss

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"sync"

	ossSDK "github.com/aliyun/aliyun-oss-go-sdk/oss"
	"jygit.jydev.jianyu360.cn/BaseService/ossService/config"
)

// bucketMap 用来缓存bucket配置，避免重复查询数据库
var bucketMap sync.Map

// accountMap 用来缓存OSS帐号信息
var accountMap sync.Map

// InitBuckets 从配置中加载bucket信息
func InitBuckets() {
	for _, bucket := range config.AppConfig.Buckets {
		bucketMap.Store(bucket.BucketID, bucket)
	}
}

// GetBucket 根据bucketID获取bucket信息，如果没有则模拟查询数据库（这里只查询一次）
func GetBucket(bucketID string) (config.BucketInfo, error) {
	if val, ok := bucketMap.Load(bucketID); ok {
		return val.(config.BucketInfo), nil
	}
	log.Println("Bucket %s not found in cache, querying database", bucketID)
	// 模拟数据库查询：实际业务中这里应查询数据库
	return config.BucketInfo{}, fmt.Errorf("bucket id %s not found", bucketID)
}

// bucketPool 缓存 bucket 对象，避免重复创建
var bucketPool sync.Map // key: bucketID, value: *ossSDK.Bucket

// GetCachedBucket 对外暴露统一入口，根据 bucketID 获取 *ossSDK.Bucket 对象
func GetCachedBucket(bucketID string) (*ossSDK.Bucket, error) {
	// 优先从对象池取
	if cached, ok := bucketPool.Load(bucketID); ok {
		return cached.(*ossSDK.Bucket), nil
	}

	// 如果不存在则依次读取 bucketInfo 和 OSS client 信息
	bucketInfo, err := GetBucket(bucketID)
	if err != nil {
		return nil, err
	}

	client, err := getOSSClient(bucketInfo.AccountID)
	if err != nil {
		return nil, err
	}

	bucket, err := client.Bucket(bucketInfo.BucketName)
	if err != nil {
		return nil, err
	}

	// 存入对象池
	bucketPool.Store(bucketID, bucket)
	return bucket, nil
}

// getOSSClient 根据OSS帐号ID从accountMap中获取OSS Client
func getOSSClient(accountID string) (*ossSDK.Client, error) {
	if val, ok := accountMap.Load(accountID); ok {
		acc := val.(config.OSSAccount)
		client, err := ossSDK.New(acc.Endpoint, acc.AccessKeyId, acc.AccessKeySecret)
		if err != nil {
			return nil, err
		}
		return client, nil
	}
	return nil, fmt.Errorf("oss account %s not found", accountID)
}

// UploadAttachment 上传附件；如果gzipEnabled为true，则进行gzip压缩（使用 io.Pipe 实现流式压缩）
func UploadAttachment(bucketID, objectName string, data io.Reader, gzipEnabled bool) error {
	bucket, err := GetCachedBucket(bucketID)
	if err != nil {
		return err
	}

	var reader io.Reader
	if gzipEnabled {
		pr, pw := io.Pipe()
		go func() {
			gzipWriter := gzip.NewWriter(pw)
			_, err := io.Copy(gzipWriter, data)
			// 注意：关闭 gzipWriter 时会将数据刷新到 pw 中
			gzipWriter.Close()
			pw.CloseWithError(err)
		}()
		reader = pr
	} else {
		reader = data
	}

	// 上传到OSS
	return bucket.PutObject(objectName, reader)
}

// DownloadAttachment 下载附件，如果检测到数据为gzip压缩，则自动解压后返回
func DownloadAttachment(bucketID, objectName string) ([]byte, error) {
	bucket, err := GetCachedBucket(bucketID)
	if err != nil {
		return nil, err
	}

	body, err := bucket.GetObject(objectName)
	if err != nil {
		return nil, err
	}
	defer body.Close()

	data, err := ioutil.ReadAll(body)
	if err != nil {
		return nil, err
	}
	// 判断是否为gzip压缩格式（判断前两个字节）
	if len(data) >= 2 && data[0] == 0x1f && data[1] == 0x8b {
		gzipReader, err := gzip.NewReader(bytes.NewReader(data))
		if err != nil {
			return nil, err
		}
		defer gzipReader.Close()
		return ioutil.ReadAll(gzipReader)
	}
	return data, nil
}

// DeleteAttachment 删除附件
func DeleteAttachment(bucketID, objectName string) error {
	bucket, err := GetCachedBucket(bucketID)
	if err != nil {
		return err
	}

	return bucket.DeleteObject(objectName)
}

// LoadOSSAccounts 加载OSS帐号信息到缓存
func LoadOSSAccounts() {
	for _, acc := range config.AppConfig.OSSAccounts {
		accountMap.Store(acc.ID, acc)
	}
}