ragflow/internal/storage/gcs.go

//
//  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()
}