mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-28 19:58:11 +08:00
Create go version storage component, but not used (#13561)
### What problem does this PR solve? Implement: minio, s3, oss, azure_sas, azure_spn, gcs, opendal ### Type of change - [x] New Feature (non-breaking change which adds functionality) --------- Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
397
internal/storage/minio.go
Normal file
397
internal/storage/minio.go
Normal file
@@ -0,0 +1,397 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/minio/minio-go/v7"
|
||||
"github.com/minio/minio-go/v7/pkg/credentials"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// MinioConfig holds MinIO storage configuration
|
||||
type MinioConfig struct {
|
||||
Host string `mapstructure:"host"` // MinIO server host (e.g., "localhost:9000")
|
||||
User string `mapstructure:"user"` // Access key
|
||||
Password string `mapstructure:"password"` // Secret key
|
||||
Secure bool `mapstructure:"secure"` // Use HTTPS
|
||||
Verify bool `mapstructure:"verify"` // Verify SSL certificates
|
||||
Bucket string `mapstructure:"bucket"` // Default bucket (optional)
|
||||
PrefixPath string `mapstructure:"prefix_path"` // Path prefix (optional)
|
||||
}
|
||||
|
||||
// MinioStorage implements Storage interface for MinIO
|
||||
type MinioStorage struct {
|
||||
client *minio.Client
|
||||
bucket string
|
||||
prefixPath string
|
||||
config *MinioConfig
|
||||
}
|
||||
|
||||
// NewMinioStorage creates a new MinIO storage instance
|
||||
func NewMinioStorage(config *MinioConfig) (*MinioStorage, error) {
|
||||
storage := &MinioStorage{
|
||||
bucket: config.Bucket,
|
||||
prefixPath: config.PrefixPath,
|
||||
config: config,
|
||||
}
|
||||
|
||||
if err := storage.connect(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return storage, nil
|
||||
}
|
||||
|
||||
func (m *MinioStorage) connect() error {
|
||||
var transport http.RoundTripper
|
||||
|
||||
// Configure transport for SSL/TLS verification
|
||||
if m.config.Secure {
|
||||
verify := m.config.Verify
|
||||
transport = &http.Transport{
|
||||
TLSClientConfig: &tls.Config{
|
||||
InsecureSkipVerify: !verify,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
client, err := minio.New(m.config.Host, &minio.Options{
|
||||
Creds: credentials.NewStaticV4(m.config.User, m.config.Password, ""),
|
||||
Secure: m.config.Secure,
|
||||
Transport: transport,
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to connect to MinIO: %w", err)
|
||||
}
|
||||
|
||||
m.client = client
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *MinioStorage) reconnect() {
|
||||
if err := m.connect(); err != nil {
|
||||
zap.L().Error("Failed to reconnect to MinIO", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MinioStorage) resolveBucketAndPath(bucket, fnm string) (string, string) {
|
||||
actualBucket := bucket
|
||||
if m.bucket != "" {
|
||||
actualBucket = m.bucket
|
||||
}
|
||||
|
||||
actualPath := fnm
|
||||
if m.bucket != "" {
|
||||
if m.prefixPath != "" {
|
||||
actualPath = fmt.Sprintf("%s/%s/%s", m.prefixPath, bucket, fnm)
|
||||
} else {
|
||||
actualPath = fmt.Sprintf("%s/%s", bucket, fnm)
|
||||
}
|
||||
} else if m.prefixPath != "" {
|
||||
actualPath = fmt.Sprintf("%s/%s", m.prefixPath, fnm)
|
||||
}
|
||||
|
||||
return actualBucket, actualPath
|
||||
}
|
||||
|
||||
// Health checks MinIO service availability
|
||||
func (m *MinioStorage) Health() bool {
|
||||
ctx := context.Background()
|
||||
|
||||
if m.bucket != "" {
|
||||
exists, err := m.client.BucketExists(ctx, m.bucket)
|
||||
if err != nil {
|
||||
zap.L().Warn("MinIO health check failed", zap.Error(err))
|
||||
return false
|
||||
}
|
||||
return exists
|
||||
}
|
||||
|
||||
_, err := m.client.ListBuckets(ctx)
|
||||
if err != nil {
|
||||
zap.L().Warn("MinIO health check failed", zap.Error(err))
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// Put uploads an object to MinIO
|
||||
func (m *MinioStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
|
||||
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 3; i++ {
|
||||
// Ensure bucket exists
|
||||
if m.bucket == "" {
|
||||
exists, err := m.client.BucketExists(ctx, bucket)
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to check bucket existence", zap.String("bucket", bucket), zap.Error(err))
|
||||
m.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
if !exists {
|
||||
if err := m.client.MakeBucket(ctx, bucket, minio.MakeBucketOptions{}); err != nil {
|
||||
zap.L().Error("Failed to create bucket", zap.String("bucket", bucket), zap.Error(err))
|
||||
m.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
reader := bytes.NewReader(binary)
|
||||
_, err := m.client.PutObject(ctx, bucket, fnm, reader, int64(len(binary)), minio.PutObjectOptions{})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to put object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
m.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("failed to put object after 3 retries")
|
||||
}
|
||||
|
||||
// Get retrieves an object from MinIO
|
||||
func (m *MinioStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
|
||||
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
obj, err := m.client.GetObject(ctx, bucket, fnm, minio.GetObjectOptions{})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to get object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
m.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
defer obj.Close()
|
||||
|
||||
buf := new(bytes.Buffer)
|
||||
if _, err := buf.ReadFrom(obj); err != nil {
|
||||
zap.L().Error("Failed to read object data", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
m.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("failed to get object after retries")
|
||||
}
|
||||
|
||||
// Rm removes an object from MinIO
|
||||
func (m *MinioStorage) Rm(bucket, fnm string, tenantID ...string) error {
|
||||
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
if err := m.client.RemoveObject(ctx, bucket, fnm, minio.RemoveObjectOptions{}); err != nil {
|
||||
zap.L().Error("Failed to remove object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ObjExist checks if an object exists in MinIO
|
||||
func (m *MinioStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
|
||||
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
exists, err := m.client.BucketExists(ctx, bucket)
|
||||
if err != nil || !exists {
|
||||
return false
|
||||
}
|
||||
|
||||
_, err = m.client.StatObject(ctx, bucket, fnm, minio.StatObjectOptions{})
|
||||
if err != nil {
|
||||
errResponse := minio.ToErrorResponse(err)
|
||||
if errResponse.Code == "NoSuchKey" || errResponse.Code == "NoSuchBucket" {
|
||||
return false
|
||||
}
|
||||
zap.L().Error("Failed to stat object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// GetPresignedURL generates a presigned URL for accessing an object
|
||||
func (m *MinioStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
bucket, fnm = m.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
url, err := m.client.PresignedGetObject(ctx, bucket, fnm, expires, nil)
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to get presigned URL", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
m.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return url.String(), nil
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("failed to get presigned URL after 10 retries")
|
||||
}
|
||||
|
||||
// BucketExists checks if a bucket exists
|
||||
func (m *MinioStorage) BucketExists(bucket string) bool {
|
||||
actualBucket := bucket
|
||||
if m.bucket != "" {
|
||||
actualBucket = m.bucket
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
exists, err := m.client.BucketExists(ctx, actualBucket)
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to check bucket existence", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return exists
|
||||
}
|
||||
|
||||
// RemoveBucket removes a bucket and all its objects
|
||||
func (m *MinioStorage) RemoveBucket(bucket string) error {
|
||||
actualBucket := bucket
|
||||
origBucket := bucket
|
||||
|
||||
if m.bucket != "" {
|
||||
actualBucket = m.bucket
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Build prefix for single-bucket mode
|
||||
prefix := ""
|
||||
if m.bucket != "" {
|
||||
if m.prefixPath != "" {
|
||||
prefix = fmt.Sprintf("%s/", m.prefixPath)
|
||||
}
|
||||
prefix += fmt.Sprintf("%s/", origBucket)
|
||||
}
|
||||
|
||||
// List and delete objects with prefix
|
||||
objectsCh := make(chan minio.ObjectInfo)
|
||||
|
||||
go func() {
|
||||
defer close(objectsCh)
|
||||
for obj := range m.client.ListObjects(ctx, actualBucket, minio.ListObjectsOptions{
|
||||
Prefix: prefix,
|
||||
Recursive: true,
|
||||
}) {
|
||||
if obj.Err != nil {
|
||||
zap.L().Error("Error listing objects", zap.Error(obj.Err))
|
||||
return
|
||||
}
|
||||
objectsCh <- obj
|
||||
}
|
||||
}()
|
||||
|
||||
for err := range m.client.RemoveObjects(ctx, actualBucket, objectsCh, minio.RemoveObjectsOptions{}) {
|
||||
zap.L().Error("Failed to remove object", zap.String("key", err.ObjectName), zap.Error(err.Err))
|
||||
}
|
||||
|
||||
// Only remove the actual bucket if not in single-bucket mode
|
||||
if m.bucket == "" {
|
||||
if err := m.client.RemoveBucket(ctx, actualBucket); err != nil {
|
||||
zap.L().Error("Failed to remove bucket", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Copy copies an object from source to destination
|
||||
func (m *MinioStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
srcBucket, srcPath = m.resolveBucketAndPath(srcBucket, srcPath)
|
||||
destBucket, destPath = m.resolveBucketAndPath(destBucket, destPath)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Ensure destination bucket exists
|
||||
if m.bucket == "" {
|
||||
exists, err := m.client.BucketExists(ctx, destBucket)
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to check bucket existence", zap.String("bucket", destBucket), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
if !exists {
|
||||
if err := m.client.MakeBucket(ctx, destBucket, minio.MakeBucketOptions{}); err != nil {
|
||||
zap.L().Error("Failed to create bucket", zap.String("bucket", destBucket), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if source object exists
|
||||
_, err := m.client.StatObject(ctx, srcBucket, srcPath, minio.StatObjectOptions{})
|
||||
if err != nil {
|
||||
zap.L().Error("Source object not found", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
// Copy object
|
||||
srcOpts := minio.CopySrcOptions{
|
||||
Bucket: srcBucket,
|
||||
Object: srcPath,
|
||||
}
|
||||
destOpts := minio.CopyDestOptions{
|
||||
Bucket: destBucket,
|
||||
Object: destPath,
|
||||
}
|
||||
|
||||
_, err = m.client.CopyObject(ctx, destOpts, srcOpts)
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to copy object", zap.String("src", fmt.Sprintf("%s/%s", srcBucket, srcPath)), zap.String("dest", fmt.Sprintf("%s/%s", destBucket, destPath)), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// Move moves an object from source to destination
|
||||
func (m *MinioStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
if m.Copy(srcBucket, srcPath, destBucket, destPath) {
|
||||
if err := m.Rm(srcBucket, srcPath); err != nil {
|
||||
zap.L().Error("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
415
internal/storage/oss.go
Normal file
415
internal/storage/oss.go
Normal file
@@ -0,0 +1,415 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/aws/smithy-go"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// OSSConfig holds Aliyun OSS storage configuration
|
||||
// OSS is compatible with S3 API
|
||||
type OSSConfig struct {
|
||||
AccessKeyID string `mapstructure:"access_key"` // OSS Access Key ID
|
||||
SecretAccessKey string `mapstructure:"secret_key"` // OSS Secret Access Key
|
||||
EndpointURL string `mapstructure:"endpoint_url"` // OSS Endpoint (e.g., "https://oss-cn-hangzhou.aliyuncs.com")
|
||||
Region string `mapstructure:"region"` // Region (e.g., "cn-hangzhou")
|
||||
Bucket string `mapstructure:"bucket"` // Default bucket (optional)
|
||||
PrefixPath string `mapstructure:"prefix_path"` // Path prefix (optional)
|
||||
SignatureVersion string `mapstructure:"signature_version"` // Signature version
|
||||
AddressingStyle string `mapstructure:"addressing_style"` // Addressing style
|
||||
}
|
||||
|
||||
// OSSStorage implements Storage interface for Aliyun OSS
|
||||
// OSS uses S3-compatible API
|
||||
type OSSStorage struct {
|
||||
client *s3.Client
|
||||
bucket string
|
||||
prefixPath string
|
||||
config *OSSConfig
|
||||
}
|
||||
|
||||
// NewOSSStorage creates a new OSS storage instance
|
||||
func NewOSSStorage(config *OSSConfig) (*OSSStorage, error) {
|
||||
storage := &OSSStorage{
|
||||
bucket: config.Bucket,
|
||||
prefixPath: config.PrefixPath,
|
||||
config: config,
|
||||
}
|
||||
|
||||
if err := storage.connect(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return storage, nil
|
||||
}
|
||||
|
||||
func (o *OSSStorage) connect() error {
|
||||
ctx := context.Background()
|
||||
|
||||
// Create static credentials
|
||||
creds := credentials.NewStaticCredentialsProvider(
|
||||
o.config.AccessKeyID,
|
||||
o.config.SecretAccessKey,
|
||||
"",
|
||||
)
|
||||
|
||||
// Load configuration
|
||||
cfg, err := config.LoadDefaultConfig(ctx,
|
||||
config.WithRegion(o.config.Region),
|
||||
config.WithCredentialsProvider(creds),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load OSS config: %w", err)
|
||||
}
|
||||
|
||||
// Create S3 client with OSS endpoint
|
||||
o.client = s3.NewFromConfig(cfg, func(opts *s3.Options) {
|
||||
opts.BaseEndpoint = aws.String(o.config.EndpointURL)
|
||||
})
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (o *OSSStorage) reconnect() {
|
||||
if err := o.connect(); err != nil {
|
||||
zap.L().Error("Failed to reconnect to OSS", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func (o *OSSStorage) resolveBucketAndPath(bucket, fnm string) (string, string) {
|
||||
actualBucket := bucket
|
||||
if o.bucket != "" {
|
||||
actualBucket = o.bucket
|
||||
}
|
||||
|
||||
actualPath := fnm
|
||||
if o.prefixPath != "" {
|
||||
actualPath = fmt.Sprintf("%s/%s", o.prefixPath, fnm)
|
||||
}
|
||||
|
||||
return actualBucket, actualPath
|
||||
}
|
||||
|
||||
// Health checks OSS service availability
|
||||
func (o *OSSStorage) Health() bool {
|
||||
bucket := o.bucket
|
||||
if bucket == "" {
|
||||
bucket = "health-check-bucket"
|
||||
}
|
||||
|
||||
fnm := "txtxtxtxt1"
|
||||
if o.prefixPath != "" {
|
||||
fnm = fmt.Sprintf("%s/%s", o.prefixPath, fnm)
|
||||
}
|
||||
binary := []byte("_t@@@1")
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Ensure bucket exists
|
||||
if !o.BucketExists(bucket) {
|
||||
_, err := o.client.CreateBucket(ctx, &s3.CreateBucketInput{
|
||||
Bucket: aws.String(bucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to create bucket for health check", zap.String("bucket", bucket), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Try to upload a test object
|
||||
reader := bytes.NewReader(binary)
|
||||
_, err := o.client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
Body: reader,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
zap.L().Error("Health check failed", zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// Put uploads an object to OSS
|
||||
func (o *OSSStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
|
||||
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
// Ensure bucket exists
|
||||
if !o.BucketExists(bucket) {
|
||||
_, err := o.client.CreateBucket(ctx, &s3.CreateBucketInput{
|
||||
Bucket: aws.String(bucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to create bucket", zap.String("bucket", bucket), zap.Error(err))
|
||||
o.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
zap.L().Info("Created bucket", zap.String("bucket", bucket))
|
||||
}
|
||||
|
||||
reader := bytes.NewReader(binary)
|
||||
_, err := o.client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
Body: reader,
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to put object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
o.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("failed to put object after retries")
|
||||
}
|
||||
|
||||
// Get retrieves an object from OSS
|
||||
func (o *OSSStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
|
||||
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
result, err := o.client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to get object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
o.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
defer result.Body.Close()
|
||||
|
||||
buf := new(bytes.Buffer)
|
||||
if _, err := buf.ReadFrom(result.Body); err != nil {
|
||||
zap.L().Error("Failed to read object data", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
o.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("failed to get object after retries")
|
||||
}
|
||||
|
||||
// Rm removes an object from OSS
|
||||
func (o *OSSStorage) Rm(bucket, fnm string, tenantID ...string) error {
|
||||
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := o.client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to remove object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ObjExist checks if an object exists in OSS
|
||||
func (o *OSSStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
|
||||
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := o.client.HeadObject(ctx, &s3.HeadObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
})
|
||||
if err != nil {
|
||||
if isOSSNotFound(err) {
|
||||
return false
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// GetPresignedURL generates a presigned URL for accessing an object
|
||||
func (o *OSSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
presignClient := s3.NewPresignClient(o.client)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
req, err := presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
}, s3.WithPresignExpires(expires))
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to generate presigned URL", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
o.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return req.URL, nil
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("failed to generate presigned URL after 10 retries")
|
||||
}
|
||||
|
||||
// BucketExists checks if a bucket exists
|
||||
func (o *OSSStorage) BucketExists(bucket string) bool {
|
||||
actualBucket := bucket
|
||||
if o.bucket != "" {
|
||||
actualBucket = o.bucket
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := o.client.HeadBucket(ctx, &s3.HeadBucketInput{
|
||||
Bucket: aws.String(actualBucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Debug("Bucket does not exist or error", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// RemoveBucket removes a bucket and all its objects
|
||||
func (o *OSSStorage) RemoveBucket(bucket string) error {
|
||||
actualBucket := bucket
|
||||
if o.bucket != "" {
|
||||
actualBucket = o.bucket
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Check if bucket exists
|
||||
if !o.BucketExists(actualBucket) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// List and delete all objects
|
||||
listInput := &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(actualBucket),
|
||||
}
|
||||
|
||||
for {
|
||||
result, err := o.client.ListObjectsV2(ctx, listInput)
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to list objects", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
for _, obj := range result.Contents {
|
||||
_, err := o.client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: aws.String(actualBucket),
|
||||
Key: obj.Key,
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to delete object", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
if result.IsTruncated == nil || !*result.IsTruncated {
|
||||
break
|
||||
}
|
||||
listInput.ContinuationToken = result.NextContinuationToken
|
||||
}
|
||||
|
||||
// Delete bucket
|
||||
_, err := o.client.DeleteBucket(ctx, &s3.DeleteBucketInput{
|
||||
Bucket: aws.String(actualBucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to delete bucket", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Copy copies an object from source to destination
|
||||
func (o *OSSStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
srcBucket, srcPath = o.resolveBucketAndPath(srcBucket, srcPath)
|
||||
destBucket, destPath = o.resolveBucketAndPath(destBucket, destPath)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
copySource := fmt.Sprintf("%s/%s", srcBucket, srcPath)
|
||||
|
||||
_, err := o.client.CopyObject(ctx, &s3.CopyObjectInput{
|
||||
Bucket: aws.String(destBucket),
|
||||
Key: aws.String(destPath),
|
||||
CopySource: aws.String(copySource),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to copy object", zap.String("src", copySource), zap.String("dest", fmt.Sprintf("%s/%s", destBucket, destPath)), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// Move moves an object from source to destination
|
||||
func (o *OSSStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
if o.Copy(srcBucket, srcPath, destBucket, destPath) {
|
||||
if err := o.Rm(srcBucket, srcPath); err != nil {
|
||||
zap.L().Error("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Helper functions
|
||||
func isOSSNotFound(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var apiErr smithy.APIError
|
||||
if errors.As(err, &apiErr) {
|
||||
return apiErr.ErrorCode() == "NotFound" || apiErr.ErrorCode() == "404" || apiErr.ErrorCode() == "NoSuchKey"
|
||||
}
|
||||
return false
|
||||
}
|
||||
425
internal/storage/s3.go
Normal file
425
internal/storage/s3.go
Normal file
@@ -0,0 +1,425 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package storage
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/aws/aws-sdk-go-v2/aws"
|
||||
"github.com/aws/aws-sdk-go-v2/config"
|
||||
"github.com/aws/aws-sdk-go-v2/credentials"
|
||||
"github.com/aws/aws-sdk-go-v2/service/s3"
|
||||
"github.com/aws/smithy-go"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// S3Config holds AWS S3 storage configuration
|
||||
type S3Config struct {
|
||||
AccessKeyID string `mapstructure:"access_key"` // AWS Access Key ID
|
||||
SecretAccessKey string `mapstructure:"secret_key"` // AWS Secret Access Key
|
||||
SessionToken string `mapstructure:"session_token"` // AWS Session Token (optional)
|
||||
Region string `mapstructure:"region_name"` // AWS Region
|
||||
EndpointURL string `mapstructure:"endpoint_url"` // Custom endpoint (optional)
|
||||
SignatureVersion string `mapstructure:"signature_version"` // Signature version
|
||||
AddressingStyle string `mapstructure:"addressing_style"` // Addressing style
|
||||
Bucket string `mapstructure:"bucket"` // Default bucket (optional)
|
||||
PrefixPath string `mapstructure:"prefix_path"` // Path prefix (optional)
|
||||
}
|
||||
|
||||
// S3Storage implements Storage interface for AWS S3
|
||||
type S3Storage struct {
|
||||
client *s3.Client
|
||||
bucket string
|
||||
prefixPath string
|
||||
config *S3Config
|
||||
}
|
||||
|
||||
// NewS3Storage creates a new S3 storage instance
|
||||
func NewS3Storage(config *S3Config) (*S3Storage, error) {
|
||||
storage := &S3Storage{
|
||||
bucket: config.Bucket,
|
||||
prefixPath: config.PrefixPath,
|
||||
config: config,
|
||||
}
|
||||
|
||||
if err := storage.connect(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return storage, nil
|
||||
}
|
||||
|
||||
func (s *S3Storage) connect() error {
|
||||
ctx := context.Background()
|
||||
|
||||
var opts []func(*config.LoadOptions) error
|
||||
|
||||
// Configure region
|
||||
if s.config.Region != "" {
|
||||
opts = append(opts, config.WithRegion(s.config.Region))
|
||||
}
|
||||
|
||||
// Configure credentials if provided
|
||||
if s.config.AccessKeyID != "" && s.config.SecretAccessKey != "" {
|
||||
creds := credentials.NewStaticCredentialsProvider(
|
||||
s.config.AccessKeyID,
|
||||
s.config.SecretAccessKey,
|
||||
s.config.SessionToken,
|
||||
)
|
||||
opts = append(opts, config.WithCredentialsProvider(creds))
|
||||
}
|
||||
|
||||
// Load configuration
|
||||
cfg, err := config.LoadDefaultConfig(ctx, opts...)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to load AWS config: %w", err)
|
||||
}
|
||||
|
||||
// Create S3 client with custom endpoint if provided
|
||||
clientOpts := []func(*s3.Options){}
|
||||
if s.config.EndpointURL != "" {
|
||||
clientOpts = append(clientOpts, func(o *s3.Options) {
|
||||
o.BaseEndpoint = aws.String(s.config.EndpointURL)
|
||||
})
|
||||
}
|
||||
|
||||
s.client = s3.NewFromConfig(cfg, clientOpts...)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *S3Storage) reconnect() {
|
||||
if err := s.connect(); err != nil {
|
||||
zap.L().Error("Failed to reconnect to S3", zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
func (s *S3Storage) resolveBucketAndPath(bucket, fnm string) (string, string) {
|
||||
actualBucket := bucket
|
||||
if s.bucket != "" {
|
||||
actualBucket = s.bucket
|
||||
}
|
||||
|
||||
actualPath := fnm
|
||||
if s.prefixPath != "" {
|
||||
actualPath = fmt.Sprintf("%s/%s/%s", s.prefixPath, bucket, fnm)
|
||||
}
|
||||
|
||||
return actualBucket, actualPath
|
||||
}
|
||||
|
||||
// Health checks S3 service availability
|
||||
func (s *S3Storage) Health() bool {
|
||||
bucket := s.bucket
|
||||
if bucket == "" {
|
||||
bucket = "health-check-bucket"
|
||||
}
|
||||
|
||||
fnm := "txtxtxtxt1"
|
||||
if s.prefixPath != "" {
|
||||
fnm = fmt.Sprintf("%s/%s", s.prefixPath, fnm)
|
||||
}
|
||||
binary := []byte("_t@@@1")
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Ensure bucket exists
|
||||
if !s.BucketExists(bucket) {
|
||||
_, err := s.client.CreateBucket(ctx, &s3.CreateBucketInput{
|
||||
Bucket: aws.String(bucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to create bucket for health check", zap.String("bucket", bucket), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Try to upload a test object
|
||||
reader := bytes.NewReader(binary)
|
||||
_, err := s.client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
Body: reader,
|
||||
})
|
||||
|
||||
if err != nil {
|
||||
zap.L().Error("Health check failed", zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// Put uploads an object to S3
|
||||
func (s *S3Storage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
|
||||
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
// Ensure bucket exists
|
||||
if !s.BucketExists(bucket) {
|
||||
_, err := s.client.CreateBucket(ctx, &s3.CreateBucketInput{
|
||||
Bucket: aws.String(bucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to create bucket", zap.String("bucket", bucket), zap.Error(err))
|
||||
s.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
zap.L().Info("Created bucket", zap.String("bucket", bucket))
|
||||
}
|
||||
|
||||
reader := bytes.NewReader(binary)
|
||||
_, err := s.client.PutObject(ctx, &s3.PutObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
Body: reader,
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to put object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
s.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
return fmt.Errorf("failed to put object after retries")
|
||||
}
|
||||
|
||||
// Get retrieves an object from S3
|
||||
func (s *S3Storage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
|
||||
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
result, err := s.client.GetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to get object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
s.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
defer result.Body.Close()
|
||||
|
||||
buf := new(bytes.Buffer)
|
||||
if _, err := buf.ReadFrom(result.Body); err != nil {
|
||||
zap.L().Error("Failed to read object data", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
s.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("failed to get object after retries")
|
||||
}
|
||||
|
||||
// Rm removes an object from S3
|
||||
func (s *S3Storage) Rm(bucket, fnm string, tenantID ...string) error {
|
||||
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to remove object", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ObjExist checks if an object exists in S3
|
||||
func (s *S3Storage) ObjExist(bucket, fnm string, tenantID ...string) bool {
|
||||
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := s.client.HeadObject(ctx, &s3.HeadObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
})
|
||||
if err != nil {
|
||||
if isS3NotFound(err) {
|
||||
return false
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// GetPresignedURL generates a presigned URL for accessing an object
|
||||
func (s *S3Storage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
bucket, fnm = s.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
presignClient := s3.NewPresignClient(s.client)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
req, err := presignClient.PresignGetObject(ctx, &s3.GetObjectInput{
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(fnm),
|
||||
}, s3.WithPresignExpires(expires))
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to generate presigned URL", zap.String("bucket", bucket), zap.String("key", fnm), zap.Error(err))
|
||||
s.reconnect()
|
||||
time.Sleep(time.Second)
|
||||
continue
|
||||
}
|
||||
|
||||
return req.URL, nil
|
||||
}
|
||||
|
||||
return "", fmt.Errorf("failed to generate presigned URL after 10 retries")
|
||||
}
|
||||
|
||||
// BucketExists checks if a bucket exists
|
||||
func (s *S3Storage) BucketExists(bucket string) bool {
|
||||
actualBucket := bucket
|
||||
if s.bucket != "" {
|
||||
actualBucket = s.bucket
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{
|
||||
Bucket: aws.String(actualBucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Debug("Bucket does not exist or error", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// RemoveBucket removes a bucket and all its objects
|
||||
func (s *S3Storage) RemoveBucket(bucket string) error {
|
||||
actualBucket := bucket
|
||||
if s.bucket != "" {
|
||||
actualBucket = s.bucket
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
// Check if bucket exists
|
||||
if !s.BucketExists(actualBucket) {
|
||||
return nil
|
||||
}
|
||||
|
||||
// List and delete all objects
|
||||
listInput := &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(actualBucket),
|
||||
}
|
||||
|
||||
for {
|
||||
result, err := s.client.ListObjectsV2(ctx, listInput)
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to list objects", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
for _, obj := range result.Contents {
|
||||
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
|
||||
Bucket: aws.String(actualBucket),
|
||||
Key: obj.Key,
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to delete object", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
if result.IsTruncated == nil || !*result.IsTruncated {
|
||||
break
|
||||
}
|
||||
listInput.ContinuationToken = result.NextContinuationToken
|
||||
}
|
||||
|
||||
// Delete bucket
|
||||
_, err := s.client.DeleteBucket(ctx, &s3.DeleteBucketInput{
|
||||
Bucket: aws.String(actualBucket),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to delete bucket", zap.String("bucket", actualBucket), zap.Error(err))
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Copy copies an object from source to destination
|
||||
func (s *S3Storage) Copy(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
srcBucket, srcPath = s.resolveBucketAndPath(srcBucket, srcPath)
|
||||
destBucket, destPath = s.resolveBucketAndPath(destBucket, destPath)
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
copySource := fmt.Sprintf("%s/%s", srcBucket, srcPath)
|
||||
|
||||
_, err := s.client.CopyObject(ctx, &s3.CopyObjectInput{
|
||||
Bucket: aws.String(destBucket),
|
||||
Key: aws.String(destPath),
|
||||
CopySource: aws.String(copySource),
|
||||
})
|
||||
if err != nil {
|
||||
zap.L().Error("Failed to copy object", zap.String("src", copySource), zap.String("dest", fmt.Sprintf("%s/%s", destBucket, destPath)), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// Move moves an object from source to destination
|
||||
func (s *S3Storage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
if s.Copy(srcBucket, srcPath, destBucket, destPath) {
|
||||
if err := s.Rm(srcBucket, srcPath); err != nil {
|
||||
zap.L().Error("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// isNotFound checks if the error is a not found error
|
||||
func isS3NotFound(err error) bool {
|
||||
if err == nil {
|
||||
return false
|
||||
}
|
||||
var apiErr smithy.APIError
|
||||
if errors.As(err, &apiErr) {
|
||||
return apiErr.ErrorCode() == "NotFound" || apiErr.ErrorCode() == "404" || apiErr.ErrorCode() == "NoSuchKey"
|
||||
}
|
||||
return false
|
||||
}
|
||||
298
internal/storage/storage_factory.go
Normal file
298
internal/storage/storage_factory.go
Normal file
@@ -0,0 +1,298 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package storage
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
"sync"
|
||||
|
||||
"github.com/spf13/viper"
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// StorageFactory creates storage instances based on configuration
|
||||
type StorageFactory struct {
|
||||
storageType StorageType
|
||||
storage Storage
|
||||
config *StorageConfig
|
||||
mu sync.RWMutex
|
||||
}
|
||||
|
||||
// StorageConfig holds all storage-related configurations
|
||||
type StorageConfig struct {
|
||||
StorageType string `mapstructure:"storage_type"`
|
||||
Minio *MinioConfig `mapstructure:"minio"`
|
||||
S3 *S3Config `mapstructure:"s3"`
|
||||
OSS *OSSConfig `mapstructure:"oss"`
|
||||
}
|
||||
|
||||
// AzureConfig holds Azure-specific configurations
|
||||
type AzureConfig struct {
|
||||
ContainerURL string `mapstructure:"container_url"`
|
||||
SASToken string `mapstructure:"sas_token"`
|
||||
AccountURL string `mapstructure:"account_url"`
|
||||
ClientID string `mapstructure:"client_id"`
|
||||
Secret string `mapstructure:"secret"`
|
||||
TenantID string `mapstructure:"tenant_id"`
|
||||
ContainerName string `mapstructure:"container_name"`
|
||||
AuthorityHost string `mapstructure:"authority_host"`
|
||||
}
|
||||
|
||||
var (
|
||||
globalFactory *StorageFactory
|
||||
once sync.Once
|
||||
)
|
||||
|
||||
// GetStorageFactory returns the singleton storage factory instance
|
||||
func GetStorageFactory() *StorageFactory {
|
||||
once.Do(func() {
|
||||
globalFactory = &StorageFactory{}
|
||||
})
|
||||
return globalFactory
|
||||
}
|
||||
|
||||
// InitStorageFactory initializes the storage factory with configuration
|
||||
func InitStorageFactory(v *viper.Viper) error {
|
||||
factory := GetStorageFactory()
|
||||
|
||||
// Get storage type from environment or config
|
||||
storageType := os.Getenv("STORAGE_IMPL")
|
||||
if storageType == "" {
|
||||
storageType = v.GetString("storage_type")
|
||||
}
|
||||
if storageType == "" {
|
||||
storageType = "MINIO" // Default storage type
|
||||
}
|
||||
|
||||
storageConfig := &StorageConfig{}
|
||||
if err := v.UnmarshalKey("storage", storageConfig); err != nil {
|
||||
return fmt.Errorf("failed to unmarshal storage config: %w", err)
|
||||
}
|
||||
storageConfig.StorageType = storageType
|
||||
|
||||
factory.config = storageConfig
|
||||
|
||||
// Initialize storage based on type
|
||||
if err := factory.initStorage(storageType, v); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
zap.L().Info("Storage factory initialized",
|
||||
zap.String("storage_type", storageType),
|
||||
)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// initStorage initializes the specific storage implementation
|
||||
func (f *StorageFactory) initStorage(storageType string, v *viper.Viper) error {
|
||||
switch storageType {
|
||||
case "MINIO":
|
||||
return f.initMinio(v)
|
||||
case "AWS_S3":
|
||||
return f.initS3(v)
|
||||
case "OSS":
|
||||
return f.initOSS(v)
|
||||
default:
|
||||
return fmt.Errorf("unsupported storage type: %s", storageType)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *StorageFactory) initMinio(v *viper.Viper) error {
|
||||
config := &MinioConfig{}
|
||||
|
||||
// Try to load from minio section first
|
||||
if v.IsSet("minio") {
|
||||
minioConfig := v.Sub("minio")
|
||||
if minioConfig != nil {
|
||||
config.Host = minioConfig.GetString("host")
|
||||
config.User = minioConfig.GetString("user")
|
||||
config.Password = minioConfig.GetString("password")
|
||||
config.Secure = minioConfig.GetBool("secure")
|
||||
config.Verify = minioConfig.GetBool("verify")
|
||||
config.Bucket = minioConfig.GetString("bucket")
|
||||
config.PrefixPath = minioConfig.GetString("prefix_path")
|
||||
}
|
||||
}
|
||||
|
||||
// Apply defaults
|
||||
if config.Host == "" {
|
||||
config.Host = "localhost:9000"
|
||||
}
|
||||
if config.User == "" {
|
||||
config.User = "minioadmin"
|
||||
}
|
||||
if config.Password == "" {
|
||||
config.Password = "minioadmin"
|
||||
}
|
||||
|
||||
storage, err := NewMinioStorage(config)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create MinIO storage: %w", err)
|
||||
}
|
||||
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.storageType = StorageMinio
|
||||
f.storage = storage
|
||||
f.config.Minio = config
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *StorageFactory) initS3(v *viper.Viper) error {
|
||||
config := &S3Config{}
|
||||
|
||||
if v.IsSet("s3") {
|
||||
s3Config := v.Sub("s3")
|
||||
if s3Config != nil {
|
||||
config.AccessKeyID = s3Config.GetString("access_key")
|
||||
config.SecretAccessKey = s3Config.GetString("secret_key")
|
||||
config.SessionToken = s3Config.GetString("session_token")
|
||||
config.Region = s3Config.GetString("region_name")
|
||||
config.EndpointURL = s3Config.GetString("endpoint_url")
|
||||
config.SignatureVersion = s3Config.GetString("signature_version")
|
||||
config.AddressingStyle = s3Config.GetString("addressing_style")
|
||||
config.Bucket = s3Config.GetString("bucket")
|
||||
config.PrefixPath = s3Config.GetString("prefix_path")
|
||||
}
|
||||
}
|
||||
|
||||
storage, err := NewS3Storage(config)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create S3 storage: %w", err)
|
||||
}
|
||||
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.storageType = StorageAWSS3
|
||||
f.storage = storage
|
||||
f.config.S3 = config
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *StorageFactory) initOSS(v *viper.Viper) error {
|
||||
config := &OSSConfig{}
|
||||
|
||||
if v.IsSet("oss") {
|
||||
ossConfig := v.Sub("oss")
|
||||
if ossConfig != nil {
|
||||
config.AccessKeyID = ossConfig.GetString("access_key")
|
||||
config.SecretAccessKey = ossConfig.GetString("secret_key")
|
||||
config.EndpointURL = ossConfig.GetString("endpoint_url")
|
||||
config.Region = ossConfig.GetString("region")
|
||||
config.Bucket = ossConfig.GetString("bucket")
|
||||
config.PrefixPath = ossConfig.GetString("prefix_path")
|
||||
config.SignatureVersion = ossConfig.GetString("signature_version")
|
||||
config.AddressingStyle = ossConfig.GetString("addressing_style")
|
||||
}
|
||||
}
|
||||
|
||||
storage, err := NewOSSStorage(config)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create OSS storage: %w", err)
|
||||
}
|
||||
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.storageType = StorageOSS
|
||||
f.storage = storage
|
||||
f.config.OSS = config
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetStorage returns the current storage instance
|
||||
func (f *StorageFactory) GetStorage() Storage {
|
||||
f.mu.RLock()
|
||||
defer f.mu.RUnlock()
|
||||
return f.storage
|
||||
}
|
||||
|
||||
// GetStorageType returns the current storage type
|
||||
func (f *StorageFactory) GetStorageType() StorageType {
|
||||
f.mu.RLock()
|
||||
defer f.mu.RUnlock()
|
||||
return f.storageType
|
||||
}
|
||||
|
||||
// Create creates a new storage instance based on the storage type
|
||||
// This is the factory method equivalent to Python's StorageFactory.create()
|
||||
func (f *StorageFactory) Create(storageType StorageType) (Storage, error) {
|
||||
var storage Storage
|
||||
var err error
|
||||
|
||||
switch storageType {
|
||||
case StorageMinio:
|
||||
if f.config.Minio != nil {
|
||||
storage, err = NewMinioStorage(f.config.Minio)
|
||||
} else {
|
||||
return nil, fmt.Errorf("MinIO config not available")
|
||||
}
|
||||
case StorageAWSS3:
|
||||
if f.config.S3 != nil {
|
||||
storage, err = NewS3Storage(f.config.S3)
|
||||
} else {
|
||||
return nil, fmt.Errorf("S3 config not available")
|
||||
}
|
||||
case StorageOSS:
|
||||
if f.config.OSS != nil {
|
||||
storage, err = NewOSSStorage(f.config.OSS)
|
||||
} else {
|
||||
return nil, fmt.Errorf("OSS config not available")
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported storage type: %v", storageType)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return storage, nil
|
||||
}
|
||||
|
||||
// SetStorage sets the storage instance (useful for testing)
|
||||
func (f *StorageFactory) SetStorage(storage Storage) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.storage = storage
|
||||
}
|
||||
|
||||
// StorageTypeMapping returns the storage type mapping (equivalent to Python's storage_mapping)
|
||||
var StorageTypeMapping = map[StorageType]func(*StorageConfig) (Storage, error){
|
||||
StorageMinio: func(config *StorageConfig) (Storage, error) {
|
||||
if config.Minio == nil {
|
||||
return nil, fmt.Errorf("MinIO config not available")
|
||||
}
|
||||
return NewMinioStorage(config.Minio)
|
||||
},
|
||||
StorageAWSS3: func(config *StorageConfig) (Storage, error) {
|
||||
if config.S3 == nil {
|
||||
return nil, fmt.Errorf("S3 config not available")
|
||||
}
|
||||
return NewS3Storage(config.S3)
|
||||
},
|
||||
StorageOSS: func(config *StorageConfig) (Storage, error) {
|
||||
if config.OSS == nil {
|
||||
return nil, fmt.Errorf("OSS config not available")
|
||||
}
|
||||
return NewOSSStorage(config.OSS)
|
||||
},
|
||||
}
|
||||
102
internal/storage/types.go
Normal file
102
internal/storage/types.go
Normal file
@@ -0,0 +1,102 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package storage
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
// ErrNotFound is returned when an object is not found
|
||||
ErrNotFound = errors.New("object not found")
|
||||
// ErrBucketNotFound is returned when a bucket is not found
|
||||
ErrBucketNotFound = errors.New("bucket not found")
|
||||
)
|
||||
|
||||
// StorageType represents the type of storage backend
|
||||
type StorageType int
|
||||
|
||||
const (
|
||||
StorageMinio StorageType = 1
|
||||
StorageAzureSpn StorageType = 2
|
||||
StorageAzureSas StorageType = 3
|
||||
StorageAWSS3 StorageType = 4
|
||||
StorageOSS StorageType = 5
|
||||
StorageOpenDAL StorageType = 6
|
||||
StorageGCS StorageType = 7
|
||||
)
|
||||
|
||||
func (s StorageType) String() string {
|
||||
switch s {
|
||||
case StorageMinio:
|
||||
return "MINIO"
|
||||
case StorageAzureSpn:
|
||||
return "AZURE_SPN"
|
||||
case StorageAzureSas:
|
||||
return "AZURE_SAS"
|
||||
case StorageAWSS3:
|
||||
return "AWS_S3"
|
||||
case StorageOSS:
|
||||
return "OSS"
|
||||
case StorageOpenDAL:
|
||||
return "OPENDAL"
|
||||
case StorageGCS:
|
||||
return "GCS"
|
||||
default:
|
||||
return "UNKNOWN"
|
||||
}
|
||||
}
|
||||
|
||||
// Storage defines the interface for storage operations
|
||||
type Storage interface {
|
||||
// Health checks the storage service availability
|
||||
Health() bool
|
||||
|
||||
// Put uploads an object to storage
|
||||
// bucket: the bucket/container name
|
||||
// fnm: the file/object name (key)
|
||||
// binary: the data to upload
|
||||
// tenantID: optional tenant identifier
|
||||
Put(bucket, fnm string, binary []byte, tenantID ...string) error
|
||||
|
||||
// Get retrieves an object from storage
|
||||
// Returns the data or nil if not found
|
||||
Get(bucket, fnm string, tenantID ...string) ([]byte, error)
|
||||
|
||||
// Rm removes an object from storage
|
||||
Rm(bucket, fnm string, tenantID ...string) error
|
||||
|
||||
// ObjExist checks if an object exists
|
||||
ObjExist(bucket, fnm string, tenantID ...string) bool
|
||||
|
||||
// GetPresignedURL generates a presigned URL for accessing an object
|
||||
// expires: duration until the URL expires
|
||||
GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error)
|
||||
|
||||
// BucketExists checks if a bucket exists
|
||||
BucketExists(bucket string) bool
|
||||
|
||||
// RemoveBucket removes a bucket and all its objects
|
||||
RemoveBucket(bucket string) error
|
||||
|
||||
// Copy copies an object from source to destination
|
||||
Copy(srcBucket, srcPath, destBucket, destPath string) bool
|
||||
|
||||
// Move moves an object from source to destination
|
||||
Move(srcBucket, srcPath, destBucket, destPath string) bool
|
||||
}
|
||||
Reference in New Issue
Block a user