mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-27 10:52:03 +08:00
Go: add more commands and GCS supports (#16741)
### Summary 1. GCS supports 2. More commands ``` RAGFlow(admin)> ping store; SUCCESS RAGFlow(admin)> ping engine; SUCCESS RAGFlow(admin)> ping cache; SUCCESS ``` --------- Signed-off-by: Jin Hai <haijin.chn@gmail.com>
This commit is contained in:
@@ -17,6 +17,7 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
@@ -26,6 +27,7 @@ import (
|
||||
"ragflow/internal/engine/redis"
|
||||
"ragflow/internal/server"
|
||||
"ragflow/internal/service"
|
||||
"ragflow/internal/storage"
|
||||
"ragflow/internal/utility"
|
||||
"strconv"
|
||||
"time"
|
||||
@@ -1250,3 +1252,38 @@ func (h *Handler) ShowModel(c *gin.Context) {
|
||||
common.SuccessWithData(c, model, "")
|
||||
|
||||
}
|
||||
|
||||
func (h *Handler) PingStore(c *gin.Context) {
|
||||
storageImpl := storage.GetStorageFactory().GetStorage()
|
||||
if storageImpl == nil {
|
||||
common.ErrorWithCode(c, int(common.CodeServerError), "storage not initialized")
|
||||
return
|
||||
}
|
||||
|
||||
if storageImpl.Health() {
|
||||
common.SuccessNoMessage(c, "SUCCESS")
|
||||
} else {
|
||||
common.ErrorWithCode(c, int(common.CodeServerError), "storage health check failed")
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) PingCache(c *gin.Context) {
|
||||
redisClient := redis.Get()
|
||||
if redisClient.Health() {
|
||||
common.SuccessNoMessage(c, "SUCCESS")
|
||||
} else {
|
||||
common.ErrorWithCode(c, int(common.CodeServerError), "cache health check failed")
|
||||
}
|
||||
}
|
||||
|
||||
func (h *Handler) PingEngine(c *gin.Context) {
|
||||
|
||||
docEngine := engine.Get()
|
||||
ctx := context.Background()
|
||||
if err := docEngine.Ping(ctx); err != nil {
|
||||
common.ErrorWithCode(c, 500, err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
common.SuccessNoMessage(c, "SUCCESS")
|
||||
}
|
||||
|
||||
@@ -102,6 +102,10 @@ func (r *Router) Setup(engine *gin.Engine) {
|
||||
queue.PUT("/messages", r.handler.PullMessageFromQueue)
|
||||
}
|
||||
|
||||
protected.GET("/store", r.handler.PingStore)
|
||||
protected.GET("/cache", r.handler.PingCache)
|
||||
protected.GET("/engine", r.handler.PingEngine)
|
||||
|
||||
protected.GET("/ingestors", r.handler.ListIngestors)
|
||||
protected.DELETE("/ingestors", r.handler.ShutdownIngestor)
|
||||
|
||||
|
||||
@@ -47,6 +47,46 @@ func (c *CLI) PingAdmin(cmd *Command) (ResponseIf, error) {
|
||||
return HandleSimpleResponse(resp, "ping admin")
|
||||
}
|
||||
|
||||
// AdminPingStoreCommand ping object store
|
||||
func (c *CLI) AdminPingStoreCommand(cmd *Command) (ResponseIf, error) {
|
||||
resp, err := c.AdminServerClient.Request("GET", "/admin/store", "admin", nil, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to ping object store: %w", err)
|
||||
}
|
||||
|
||||
return HandleSimpleResponse(resp, fmt.Sprintf("ping object store"))
|
||||
}
|
||||
|
||||
// AdminPingEngineCommand ping document engine
|
||||
func (c *CLI) AdminPingEngineCommand(cmd *Command) (ResponseIf, error) {
|
||||
resp, err := c.AdminServerClient.Request("GET", "/admin/engine", "admin", nil, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to ping document engine: %w", err)
|
||||
}
|
||||
|
||||
return HandleSimpleResponse(resp, fmt.Sprintf("ping document engine"))
|
||||
}
|
||||
|
||||
// AdminPingMQCommand ping message queue
|
||||
func (c *CLI) AdminPingMQCommand(cmd *Command) (ResponseIf, error) {
|
||||
resp, err := c.AdminServerClient.Request("GET", "/admin/queue", "admin", nil, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to ping message queue: %w", err)
|
||||
}
|
||||
|
||||
return HandleSimpleResponse(resp, fmt.Sprintf("ping message queue"))
|
||||
}
|
||||
|
||||
// AdminPingCacheCommand ping cache
|
||||
func (c *CLI) AdminPingCacheCommand(cmd *Command) (ResponseIf, error) {
|
||||
resp, err := c.AdminServerClient.Request("GET", "/admin/cache", "web", nil, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to ping cache: %w", err)
|
||||
}
|
||||
|
||||
return HandleSimpleResponse(resp, fmt.Sprintf("ping cache"))
|
||||
}
|
||||
|
||||
// AdminShowVersionCommand show RAGFlow admin version
|
||||
func (c *CLI) AdminShowVersionCommand(cmd *Command) (ResponseIf, error) {
|
||||
resp, err := c.AdminServerClient.Request("GET", "/admin/version", "web", nil, nil)
|
||||
@@ -2633,3 +2673,24 @@ func (c *CLI) AdminShowLogLevelCommand(cmd *Command) (ResponseIf, error) {
|
||||
|
||||
return HandleCommonDataResponse(resp, fmt.Sprintf("get log level config"))
|
||||
}
|
||||
|
||||
func (c *CLI) AdminListBucketObjects(cmd *Command) (ResponseIf, error) {
|
||||
|
||||
if c.Config.CLIMode != AdminMode {
|
||||
return nil, fmt.Errorf("this command is only allowed in ADMIN mode or already login")
|
||||
}
|
||||
|
||||
bucketID, ok := cmd.Params["bucket_id"].(string)
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("bucket_id not provided")
|
||||
}
|
||||
|
||||
endPoint := fmt.Sprintf("/admin/bucket/%s/objects", bucketID)
|
||||
|
||||
resp, err := c.AdminServerClient.Request("GET", endPoint, "web", nil, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get bucket objects: %w", err)
|
||||
}
|
||||
|
||||
return HandleCommonResponse(resp, fmt.Sprintf("get bucket objects"))
|
||||
}
|
||||
|
||||
@@ -70,12 +70,30 @@ func (p *Parser) parseAdminLogout() (*Command, error) {
|
||||
}
|
||||
|
||||
func (p *Parser) parseAdminPingServer() (*Command, error) {
|
||||
cmd := NewCommand("admin_ping_server")
|
||||
p.nextToken()
|
||||
// Semicolon is optional
|
||||
if p.curToken.Type == TokenSemicolon {
|
||||
p.nextToken() // consume PING
|
||||
|
||||
var cmd *Command
|
||||
|
||||
switch p.curToken.Type {
|
||||
case TokenStore:
|
||||
p.nextToken()
|
||||
cmd = NewCommand("admin_ping_store")
|
||||
case TokenEngine:
|
||||
p.nextToken()
|
||||
cmd = NewCommand("admin_ping_engine")
|
||||
case TokenMQ:
|
||||
p.nextToken()
|
||||
cmd = NewCommand("admin_ping_mq")
|
||||
case TokenCache:
|
||||
p.nextToken()
|
||||
cmd = NewCommand("admin_ping_cache")
|
||||
case TokenSemicolon, TokenEOF:
|
||||
p.nextToken()
|
||||
cmd = NewCommand("admin_ping_server")
|
||||
default:
|
||||
return nil, fmt.Errorf("expected semicolon after PING")
|
||||
}
|
||||
|
||||
return cmd, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -41,8 +41,16 @@ func (c *CLI) ExecuteAdminCommand(cmd *Command) (ResponseIf, error) {
|
||||
return c.LoginUserByCommand(cmd)
|
||||
case "admin_logout":
|
||||
return c.Logout()
|
||||
case "admin_ping_store":
|
||||
return c.AdminPingStoreCommand(cmd)
|
||||
case "admin_ping_engine":
|
||||
return c.AdminPingEngineCommand(cmd)
|
||||
case "admin_ping_mq":
|
||||
return c.AdminPingMQCommand(cmd)
|
||||
case "admin_ping_cache":
|
||||
return c.AdminPingCacheCommand(cmd)
|
||||
case "admin_ping_server":
|
||||
return c.PingByCommand(cmd)
|
||||
return c.PingServerByCommand(cmd)
|
||||
case "benchmark":
|
||||
return c.RunBenchmark(cmd)
|
||||
case "admin_list_services":
|
||||
@@ -281,6 +289,8 @@ func (c *CLI) ExecuteAdminCommand(cmd *Command) (ResponseIf, error) {
|
||||
return c.CommonUseAPIServerCommand(cmd)
|
||||
case "admin_use_admin_server":
|
||||
return c.CommonUseAdminServerCommand(cmd)
|
||||
case "admin_list_bucket_objects":
|
||||
return c.AdminListBucketObjects(cmd)
|
||||
default:
|
||||
return nil, fmt.Errorf("command '%s' would be executed with API", cmd.Type)
|
||||
}
|
||||
@@ -294,7 +304,7 @@ func (c *CLI) ExecuteUserCommand(cmd *Command) (ResponseIf, error) {
|
||||
case "api_logout":
|
||||
return c.Logout()
|
||||
case "api_ping_server":
|
||||
return c.PingByCommand(cmd)
|
||||
return c.PingServerByCommand(cmd)
|
||||
case "api_list_configs":
|
||||
return c.ListConfigs(cmd)
|
||||
case "api_set_log_level":
|
||||
|
||||
@@ -109,7 +109,7 @@ func (c *CLI) LoginUserInteractive(email, password string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *CLI) PingByCommand(cmd *Command) (ResponseIf, error) {
|
||||
func (c *CLI) PingServerByCommand(cmd *Command) (ResponseIf, error) {
|
||||
iterations := 1
|
||||
if iterationsParam, ok := cmd.Params["iterations"]; ok {
|
||||
iterations = int(iterationsParam.(float64))
|
||||
|
||||
@@ -335,6 +335,8 @@ func (l *Lexer) lookupIdent(ident string) Token {
|
||||
return Token{Type: TokenMax, Value: ident}
|
||||
case "STORE":
|
||||
return Token{Type: TokenStore, Value: ident}
|
||||
case "ENGINE":
|
||||
return Token{Type: TokenEngine, Value: ident}
|
||||
case "STREAM":
|
||||
return Token{Type: TokenStream, Value: ident}
|
||||
case "LS":
|
||||
@@ -491,6 +493,8 @@ func (l *Lexer) lookupIdent(ident string) Token {
|
||||
return Token{Type: TokenIngestors, Value: ident}
|
||||
case "INGESTION":
|
||||
return Token{Type: TokenIngestion, Value: ident}
|
||||
case "CACHE":
|
||||
return Token{Type: TokenCache, Value: ident}
|
||||
case "MQ":
|
||||
return Token{Type: TokenMQ, Value: ident}
|
||||
case "PUBLISH":
|
||||
|
||||
@@ -135,6 +135,7 @@ const (
|
||||
TokenVector
|
||||
TokenSize
|
||||
TokenStore
|
||||
TokenEngine
|
||||
TokenName // For ALTER PROVIDER <name> NAME <new_name>
|
||||
TokenBalance
|
||||
TokenInstance
|
||||
@@ -176,6 +177,7 @@ const (
|
||||
TokenStart
|
||||
TokenStop
|
||||
TokenIngestion
|
||||
TokenCache
|
||||
TokenMQ
|
||||
TokenPublish
|
||||
TokenPull
|
||||
|
||||
@@ -19,6 +19,7 @@ package elasticsearch
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
@@ -59,6 +60,7 @@ func NewEngine(cfg interface{}) (*elasticsearchEngine, error) {
|
||||
Transport: &http.Transport{
|
||||
MaxIdleConnsPerHost: 10,
|
||||
ResponseHeaderTimeout: 30 * time.Second,
|
||||
TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
|
||||
@@ -123,18 +123,18 @@ func (n *NatsEngine) ShowMessageQueue() (map[string]string, error) {
|
||||
result["message_count"] = strconv.FormatUint(info.State.Msgs, 10)
|
||||
|
||||
consumer, err := n.stream.Consumer(ctx, "RAGFLOW_CONSUMER")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get existing consumer: %w", err)
|
||||
if err == nil {
|
||||
var consumerInfo *jetstream.ConsumerInfo
|
||||
consumerInfo, err = consumer.Info(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get consumer info: %w", err)
|
||||
}
|
||||
result["pending_count"] = strconv.FormatUint(consumerInfo.NumPending, 10)
|
||||
result["waiting_count"] = strconv.Itoa(consumerInfo.NumWaiting)
|
||||
result["ack_pending_count"] = strconv.Itoa(consumerInfo.NumAckPending)
|
||||
result["redelivered_count"] = strconv.Itoa(consumerInfo.NumRedelivered)
|
||||
}
|
||||
|
||||
consumerInfo, err := consumer.Info(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get consumer info: %w", err)
|
||||
}
|
||||
result["pending_count"] = strconv.FormatUint(consumerInfo.NumPending, 10)
|
||||
result["waiting_count"] = strconv.Itoa(consumerInfo.NumWaiting)
|
||||
result["ack_pending_count"] = strconv.Itoa(consumerInfo.NumAckPending)
|
||||
result["redelivered_count"] = strconv.Itoa(consumerInfo.NumRedelivered)
|
||||
return result, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -210,12 +210,14 @@ type StorageConfig struct {
|
||||
Minio *MinioConfig `mapstructure:"minio"`
|
||||
S3 *S3Config `mapstructure:"s3"`
|
||||
OSS *OSSConfig `mapstructure:"oss"`
|
||||
GCS *GCSConfig `mapstructure:"gcs"`
|
||||
}
|
||||
|
||||
const (
|
||||
StorageOSS StorageType = "oss"
|
||||
StorageS3 StorageType = "s3"
|
||||
StorageMinio StorageType = "minio"
|
||||
StorageGCS StorageType = "gcs"
|
||||
)
|
||||
|
||||
// OSSConfig holds Aliyun OSS storage configuration
|
||||
@@ -231,6 +233,12 @@ type OSSConfig struct {
|
||||
AddressingStyle string `mapstructure:"addressing_style"` // Addressing style
|
||||
}
|
||||
|
||||
type GCSConfig struct {
|
||||
Bucket string `mapstructure:"bucket"` // Default bucket (optional)
|
||||
PrefixPath string `mapstructure:"prefix_path"` // Path prefix (optional)
|
||||
EndpointURL string `mapstructure:"endpoint_url"` // Custom endpoint (optional)
|
||||
}
|
||||
|
||||
// MinioConfig holds MinIO storage configuration
|
||||
type MinioConfig struct {
|
||||
Host string `mapstructure:"host"` // MinIO server host (e.g., "localhost:9000")
|
||||
@@ -534,6 +542,8 @@ func FromEnvironments() error {
|
||||
globalConfig.StorageEngine.Type = StorageS3
|
||||
case "oss":
|
||||
globalConfig.StorageEngine.Type = StorageOSS
|
||||
case "gcs":
|
||||
globalConfig.StorageEngine.Type = StorageGCS
|
||||
case "":
|
||||
// Default
|
||||
if globalConfig.StorageEngine.Type == "" {
|
||||
@@ -788,6 +798,19 @@ func FromConfigFile(configPath string) error {
|
||||
}
|
||||
}
|
||||
|
||||
if v.IsSet("gcs") {
|
||||
gcsConfig := v.Sub("gcs")
|
||||
if gcsConfig != nil {
|
||||
if globalConfig.StorageEngine.GCS == nil {
|
||||
globalConfig.StorageEngine.GCS = &GCSConfig{
|
||||
Bucket: gcsConfig.GetString("bucket"),
|
||||
PrefixPath: gcsConfig.GetString("prefix_path"),
|
||||
EndpointURL: gcsConfig.GetString("endpoint_url"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if v.IsSet("minio_0") {
|
||||
minioConfig := v.Sub("minio_0")
|
||||
if minioConfig != nil {
|
||||
|
||||
@@ -1226,6 +1226,11 @@ func (s *chunkImageStorage) Get(bucket, fnm string, tenantID ...string) ([]byte,
|
||||
}
|
||||
func (s *chunkImageStorage) Remove(bucket, fnm string, tenantID ...string) error { return nil }
|
||||
func (s *chunkImageStorage) ObjExist(bucket, fnm string, tenantID ...string) bool { return s.exists }
|
||||
|
||||
// ListObjects lists all objects in a bucket
|
||||
func (s *chunkImageStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
return []string{}, nil
|
||||
}
|
||||
func (s *chunkImageStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
@@ -1233,6 +1238,7 @@ func (s *chunkImageStorage) BucketExists(bucket string) bool
|
||||
func (s *chunkImageStorage) RemoveBucket(bucket string) error { return nil }
|
||||
func (s *chunkImageStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool { return false }
|
||||
func (s *chunkImageStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool { return false }
|
||||
func (s *chunkImageStorage) Close() error { return nil }
|
||||
|
||||
func mustEncodePNG(t *testing.T, rect image.Rectangle) []byte {
|
||||
t.Helper()
|
||||
|
||||
@@ -70,6 +70,9 @@ func (f *fakeUploadStorage) ObjExist(bucket, fnm string, tenantID ...string) boo
|
||||
_, ok := f.objects[f.key(bucket, fnm)]
|
||||
return ok
|
||||
}
|
||||
func (f *fakeUploadStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
return []string{}, nil
|
||||
}
|
||||
func (f *fakeUploadStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
return "", nil
|
||||
}
|
||||
@@ -90,6 +93,7 @@ func (f *fakeUploadStorage) Move(srcBucket, srcPath, destBucket, destPath string
|
||||
delete(f.objects, f.key(srcBucket, srcPath))
|
||||
return true
|
||||
}
|
||||
func (f *fakeUploadStorage) Close() error { return nil }
|
||||
|
||||
type fakeChatDocEngine struct{}
|
||||
|
||||
|
||||
@@ -72,6 +72,10 @@ func (f *fakeStorage) BucketExists(bucket string) bool {
|
||||
panic("not implemented in fakeStorage")
|
||||
}
|
||||
|
||||
func (f *fakeStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
panic("not implemented in fakeStorage")
|
||||
}
|
||||
|
||||
func (f *fakeStorage) RemoveBucket(bucket string) error {
|
||||
panic("not implemented in fakeStorage")
|
||||
}
|
||||
@@ -84,6 +88,8 @@ func (f *fakeStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool
|
||||
panic("not implemented in fakeStorage")
|
||||
}
|
||||
|
||||
func (f *fakeStorage) Close() error { return nil }
|
||||
|
||||
func setupFileContentPermissionDB(t *testing.T, accessible bool) {
|
||||
t.Helper()
|
||||
|
||||
|
||||
253
internal/storage/gcs.go
Normal file
253
internal/storage/gcs.go
Normal file
@@ -0,0 +1,253 @@
|
||||
//
|
||||
// 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 (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/server"
|
||||
"time"
|
||||
|
||||
"cloud.google.com/go/storage"
|
||||
"go.uber.org/zap"
|
||||
"google.golang.org/api/iterator"
|
||||
)
|
||||
|
||||
// GCSStorage implements Storage interface for GCS
|
||||
type GCSStorage struct {
|
||||
client *storage.Client
|
||||
config *server.GCSConfig
|
||||
}
|
||||
|
||||
// NewGCSStorage creates a new GCS storage instance
|
||||
func NewGCSStorage(config *server.GCSConfig) (*GCSStorage, error) {
|
||||
gcsStorage := &GCSStorage{
|
||||
config: config,
|
||||
}
|
||||
|
||||
if err := gcsStorage.connect(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return gcsStorage, nil
|
||||
}
|
||||
|
||||
func (m *GCSStorage) connect() error {
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
client, err := storage.NewClient(ctx)
|
||||
if err != nil {
|
||||
common.Fatal(fmt.Sprintf("Failed to create client: %s", err.Error()))
|
||||
}
|
||||
|
||||
m.client = client
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *GCSStorage) reconnect() {
|
||||
if err := m.connect(); err != nil {
|
||||
common.Fatal(fmt.Sprintf("Failed to reconnect to GCS, %s", err.Error()))
|
||||
}
|
||||
}
|
||||
|
||||
// Health checks GCS service availability
|
||||
func (m *GCSStorage) Health() bool {
|
||||
return m.BucketExists(m.config.Bucket)
|
||||
}
|
||||
|
||||
// Put uploads an object to GCS
|
||||
func (m *GCSStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error {
|
||||
ctx := context.Background()
|
||||
|
||||
obj := m.client.Bucket(bucket).Object(fnm)
|
||||
w := obj.NewWriter(ctx)
|
||||
|
||||
if _, err := w.Write(binary); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := w.Close(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Get retrieves an object from GCS
|
||||
func (m *GCSStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) {
|
||||
ctx := context.Background()
|
||||
|
||||
r, err := m.client.Bucket(bucket).Object(fnm).NewReader(ctx)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer r.Close()
|
||||
|
||||
data, err := io.ReadAll(r)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return data, nil
|
||||
}
|
||||
|
||||
// Remove removes an object from GCS
|
||||
func (m *GCSStorage) Remove(bucketName, objectName string, tenantID ...string) error {
|
||||
ctx := context.Background()
|
||||
|
||||
obj := m.client.Bucket(bucketName).Object(objectName)
|
||||
if err := obj.Delete(ctx); err != nil {
|
||||
return fmt.Errorf("fail to delete object: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// ObjExist checks if an object exists in GCS
|
||||
func (m *GCSStorage) ObjExist(bucketName, objectName string, tenantID ...string) bool {
|
||||
ctx := context.Background()
|
||||
|
||||
obj := m.client.Bucket(bucketName).Object(objectName)
|
||||
|
||||
_, err := obj.Attrs(ctx)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
func (m *GCSStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
ctx := context.Background()
|
||||
|
||||
bucketObject := m.client.Bucket(bucket)
|
||||
it := bucketObject.Objects(ctx, nil)
|
||||
|
||||
var objects []string
|
||||
for {
|
||||
attrs, err := it.Next()
|
||||
if errors.Is(err, iterator.Done) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
objects = append(objects, attrs.Name)
|
||||
}
|
||||
|
||||
return objects, nil
|
||||
}
|
||||
|
||||
// GetPresignedURL generates a presigned URL for accessing an object
|
||||
func (m *GCSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
|
||||
bucketObject := m.client.Bucket(bucket)
|
||||
objectPath := fmt.Sprintf("%s/%s", bucket, fnm)
|
||||
opts := &storage.SignedURLOptions{
|
||||
Method: "GET",
|
||||
Expires: time.Now().Add(expires * time.Second),
|
||||
}
|
||||
url, err := bucketObject.SignedURL(objectPath, opts)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
return url, nil
|
||||
}
|
||||
|
||||
// BucketExists checks if a bucket exists
|
||||
func (m *GCSStorage) BucketExists(bucket string) bool {
|
||||
actualBucket := bucket
|
||||
if m.config.Bucket != "" {
|
||||
actualBucket = m.config.Bucket
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
_, err := m.client.Bucket(actualBucket).Attrs(ctx)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// RemoveBucket removes a bucket and all its objects
|
||||
func (m *GCSStorage) RemoveBucket(bucketName string) error {
|
||||
if bucketName == "" {
|
||||
return fmt.Errorf("attempt to delete bucket without name")
|
||||
}
|
||||
|
||||
ctx := context.Background()
|
||||
|
||||
bucket := m.client.Bucket(bucketName)
|
||||
|
||||
it := bucket.Objects(ctx, nil)
|
||||
for {
|
||||
attrs, err := it.Next()
|
||||
if errors.Is(err, iterator.Done) {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err = bucket.Object(attrs.Name).Delete(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if err := bucket.Delete(ctx); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// Copy copies an object from source to destination
|
||||
func (m *GCSStorage) Copy(srcBucket, srcObject, destBucket, destObject string) bool {
|
||||
ctx := context.Background()
|
||||
src := m.client.Bucket(srcBucket).Object(srcObject)
|
||||
dst := m.client.Bucket(destBucket).Object(destObject)
|
||||
copier := dst.CopierFrom(src)
|
||||
_, err := copier.Run(ctx)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// Move moves an object from source to destination
|
||||
func (m *GCSStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
if m.Copy(srcBucket, srcPath, destBucket, destPath) {
|
||||
if err := m.Remove(srcBucket, srcPath); err != nil {
|
||||
common.Warn("Failed to remove source object after copy", zap.String("bucket", srcBucket), zap.String("key", srcPath), zap.Error(err))
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (m *GCSStorage) Close() error {
|
||||
common.Info("Closing GCS client")
|
||||
return m.client.Close()
|
||||
}
|
||||
@@ -123,6 +123,25 @@ func (m *MemoryStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
|
||||
return ok
|
||||
}
|
||||
|
||||
func (m *MemoryStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
m.mu.RLock()
|
||||
defer m.mu.RUnlock()
|
||||
|
||||
bucketMap, ok := m.objects[bucket]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("memory storage: bucket %q: %w", bucket, ErrMemoryNotFound)
|
||||
}
|
||||
|
||||
var objects []string
|
||||
for k := range bucketMap {
|
||||
objects = append(objects, k)
|
||||
}
|
||||
|
||||
return objects, nil
|
||||
}
|
||||
|
||||
// ListObjects lists all objects in a bucket
|
||||
|
||||
// GetPresignedURL returns a deterministic, non-network URL string for tests.
|
||||
// Format: memory://<bucket>/<key>?exp=<unix-seconds>
|
||||
func (m *MemoryStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
@@ -202,6 +221,8 @@ func (m *MemoryStorage) Move(srcBucket, srcPath, destBucket, destPath string) bo
|
||||
return true
|
||||
}
|
||||
|
||||
func (m *MemoryStorage) Close() error { return nil }
|
||||
|
||||
// Inspect returns a stable snapshot of all (bucket, key, size) entries
|
||||
// currently held by the backend. Intended for test diagnostics and
|
||||
// cleanup assertions. The slice is freshly allocated and safe to mutate.
|
||||
|
||||
@@ -273,6 +273,24 @@ func (m *MinioStorage) BucketExists(bucket string) bool {
|
||||
return exists
|
||||
}
|
||||
|
||||
func (m *MinioStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
ctx := context.Background()
|
||||
|
||||
var objects []string
|
||||
for obj := range m.client.ListObjects(ctx, bucket, minio.ListObjectsOptions{
|
||||
Prefix: "",
|
||||
Recursive: true,
|
||||
}) {
|
||||
if obj.Err != nil {
|
||||
common.Warn("Failed to list objects", zap.Error(obj.Err))
|
||||
return nil, obj.Err
|
||||
}
|
||||
objects = append(objects, obj.Key)
|
||||
}
|
||||
|
||||
return objects, nil
|
||||
}
|
||||
|
||||
// RemoveBucket removes a bucket and all its objects
|
||||
func (m *MinioStorage) RemoveBucket(bucket string) error {
|
||||
actualBucket := bucket
|
||||
@@ -384,3 +402,5 @@ func (m *MinioStorage) Move(srcBucket, srcPath, destBucket, destPath string) boo
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func (m *MinioStorage) Close() error { return nil }
|
||||
|
||||
@@ -256,7 +256,27 @@ func (o *OSSStorage) ObjExist(bucket, fnm string, tenantID ...string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// GetPresignedURL generates a presigned URL for accessing an object
|
||||
func (o *OSSStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
ctx := context.Background()
|
||||
|
||||
listInput := &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(bucket),
|
||||
}
|
||||
|
||||
result, err := o.client.ListObjectsV2(ctx, listInput)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var objects []string
|
||||
|
||||
for _, obj := range result.Contents {
|
||||
objects = append(objects, *obj.Key)
|
||||
}
|
||||
|
||||
return objects, nil
|
||||
}
|
||||
|
||||
func (o *OSSStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) {
|
||||
bucket, fnm = o.resolveBucketAndPath(bucket, fnm)
|
||||
|
||||
@@ -390,6 +410,8 @@ func (o *OSSStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool
|
||||
return false
|
||||
}
|
||||
|
||||
func (o *OSSStorage) Close() error { return nil }
|
||||
|
||||
// Helper functions
|
||||
func isOSSNotFound(err error) bool {
|
||||
if err == nil {
|
||||
|
||||
@@ -264,6 +264,27 @@ func (s *S3Storage) ObjExist(bucket, fnm string, tenantID ...string) bool {
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *S3Storage) ListObjects(bucket string, tenantID ...string) ([]string, error) {
|
||||
ctx := context.Background()
|
||||
|
||||
listInput := &s3.ListObjectsV2Input{
|
||||
Bucket: aws.String(bucket),
|
||||
}
|
||||
|
||||
result, err := s.client.ListObjectsV2(ctx, listInput)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var objects []string
|
||||
|
||||
for _, obj := range result.Contents {
|
||||
objects = append(objects, *obj.Key)
|
||||
}
|
||||
|
||||
return objects, nil
|
||||
}
|
||||
|
||||
// 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)
|
||||
@@ -398,6 +419,8 @@ func (s *S3Storage) Move(srcBucket, srcPath, destBucket, destPath string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
func (s *S3Storage) Close() error { return nil }
|
||||
|
||||
// isNotFound checks if the error is a not found error
|
||||
func isS3NotFound(err error) bool {
|
||||
if err == nil {
|
||||
|
||||
@@ -60,7 +60,14 @@ func InitStorageFactory() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// initStorage initializes the specific storage implementation
|
||||
// CloseStorage closes the storage connection
|
||||
func CloseStorage() error {
|
||||
factory := GetStorageFactory()
|
||||
factory.mu.Lock()
|
||||
defer factory.mu.Unlock()
|
||||
return factory.storage.Close()
|
||||
}
|
||||
|
||||
func (f *StorageFactory) initStorage() error {
|
||||
switch f.config.Type {
|
||||
case "minio":
|
||||
@@ -69,6 +76,8 @@ func (f *StorageFactory) initStorage() error {
|
||||
return f.initS3(f.config.S3)
|
||||
case "oss":
|
||||
return f.initOSS(f.config.OSS)
|
||||
case "gcs":
|
||||
return f.initGCS(f.config.GCS)
|
||||
default:
|
||||
return fmt.Errorf("unsupported storage type: %s", f.config.Type)
|
||||
}
|
||||
@@ -120,6 +129,22 @@ func (f *StorageFactory) initOSS(ossConfig *server.OSSConfig) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *StorageFactory) initGCS(gcsConfig *server.GCSConfig) error {
|
||||
|
||||
storage, err := NewGCSStorage(gcsConfig)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create GCS storage: %w", err)
|
||||
}
|
||||
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.storageType = StorageGCS
|
||||
f.storage = storage
|
||||
f.config.GCS = gcsConfig
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetStorage returns the current storage instance
|
||||
func (f *StorageFactory) GetStorage() Storage {
|
||||
f.mu.RLock()
|
||||
|
||||
@@ -76,6 +76,9 @@ type Storage interface {
|
||||
// ObjExist checks if an object exists
|
||||
ObjExist(bucket, fnm string, tenantID ...string) bool
|
||||
|
||||
// ListObjects list all objects of the bucket
|
||||
ListObjects(bucket string, tenantID ...string) ([]string, error)
|
||||
|
||||
// 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)
|
||||
@@ -91,4 +94,7 @@ type Storage interface {
|
||||
|
||||
// Move moves an object from source to destination
|
||||
Move(srcBucket, srcPath, destBucket, destPath string) bool
|
||||
|
||||
// Close closes the storage connection
|
||||
Close() error
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user