ai-gateway-go/internal/workbench/files.go

package workbench

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"strings"
	"time"

	"aigateway.local/core/internal/platform/storage"
	"github.com/jackc/pgx/v5"
)

// FileObject is the metadata row for one object stored in MinIO. The object
// body lives in the bucket; this table is the searchable index plus the access
// control (personal scope is bound to a portal user, system scope to admins).
type FileObject struct {
	ID            string    `json:"id"`
	ObjectKey     string    `json:"-"`
	OriginalName  string    `json:"original_name"`
	ContentType   string    `json:"content_type"`
	SizeBytes     int64     `json:"size_bytes"`
	ContentSHA256 string    `json:"content_sha256"`
	Scope         string    `json:"scope"`
	OwnerUserID   *string   `json:"owner_user_id,omitempty"`
	CreatedBy     string    `json:"created_by"`
	CreatedAt     time.Time `json:"created_at"`
	UpdatedAt     time.Time `json:"updated_at"`
}

// FileService stores object metadata in PostgreSQL and the body in MinIO. Both
// writes are kept consistent: the object is uploaded first and only then is the
// metadata row inserted; any failure rolls the object back.
type FileService struct {
	assets *Service
	store  *storage.Client
}

func NewFileService(assets *Service, store *storage.Client) *FileService {
	return &FileService{assets: assets, store: store}
}

func (s *FileService) MaxBytes() int64 { return s.store.MaxBytes() }

// countingReader tallies how many bytes flow through it so uploads can enforce
// the configured size ceiling without buffering the whole file in memory.
type countingReader struct {
	r io.Reader
	n *int64
}

func (c *countingReader) Read(p []byte) (int, error) {
	n, err := c.r.Read(p)
	*c.n += int64(n)
	return n, err
}

// Upload streams the body into the object store and records its metadata. The
// owner must be set for the personal scope and nil for the system scope.
func (s *FileService) Upload(ctx context.Context, scope string, ownerUserID *string, createdBy, originalName, contentType string, body io.Reader) (FileObject, error) {
	originalName = strings.TrimSpace(originalName)
	if originalName == "" || len(originalName) > 255 {
		return FileObject{}, errors.New("文件名无效")
	}
	if scope != "personal" && scope != "system" {
		return FileObject{}, errors.New("无效的文件范围")
	}
	if contentType == "" {
		contentType = "application/octet-stream"
	}
	objectKey, err := newUUID()
	if err != nil {
		return FileObject{}, err
	}
	objectKey = scope + "/" + objectKey
	var size int64
	hasher := sha256.New()
	limited := io.LimitReader(body, s.store.MaxBytes()+1)
	counted := &countingReader{r: limited, n: &size}
	rollbackObject := func() { _ = s.store.DeleteObject(ctx, objectKey) }
	if _, err = s.store.PutObject(ctx, objectKey, io.TeeReader(counted, hasher), -1, contentType); err != nil {
		return FileObject{}, err
	}
	if size > s.store.MaxBytes() {
		rollbackObject()
		return FileObject{}, errors.New("文件超过大小上限")
	}
	id, err := newUUID()
	if err != nil {
		rollbackObject()
		return FileObject{}, err
	}
	var owner any
	if ownerUserID != nil && strings.TrimSpace(*ownerUserID) != "" {
		owner = *ownerUserID
	}
	var obj FileObject
	err = s.assets.pool.QueryRow(ctx, `INSERT INTO gateway.file_objects(id,object_key,original_name,content_type,size_bytes,content_sha256,scope,owner_user_id,created_by) VALUES($1,$2,$3,$4,$5,$6,$7,$8,$9) RETURNING id::text,object_key,original_name,content_type,size_bytes,content_sha256,scope,owner_user_id::text,created_by,created_at,updated_at`, id, objectKey, originalName, contentType, size, hex.EncodeToString(hasher.Sum(nil)), scope, owner, createdBy).Scan(&obj.ID, &obj.ObjectKey, &obj.OriginalName, &obj.ContentType, &obj.SizeBytes, &obj.ContentSHA256, &obj.Scope, &obj.OwnerUserID, &obj.CreatedBy, &obj.CreatedAt, &obj.UpdatedAt)
	if err != nil {
		rollbackObject()
		return FileObject{}, err
	}
	return obj, nil
}

const fileObjectSelect = `SELECT id::text,object_key,original_name,content_type,size_bytes,content_sha256,scope,owner_user_id::text,created_by,created_at,updated_at FROM gateway.file_objects`

func scanFileObject(row pgx.Row) (FileObject, error) {
	var obj FileObject
	err := row.Scan(&obj.ID, &obj.ObjectKey, &obj.OriginalName, &obj.ContentType, &obj.SizeBytes, &obj.ContentSHA256, &obj.Scope, &obj.OwnerUserID, &obj.CreatedBy, &obj.CreatedAt, &obj.UpdatedAt)
	return obj, mapNotFound(err)
}

func (s *FileService) ListPersonal(ctx context.Context, ownerUserID string) ([]FileObject, error) {
	rows, err := s.assets.pool.Query(ctx, fileObjectSelect+` WHERE scope='personal' AND owner_user_id=$1 ORDER BY created_at DESC`, ownerUserID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []FileObject{}
	for rows.Next() {
		obj, err := scanFileObject(rows)
		if err != nil {
			return nil, err
		}
		items = append(items, obj)
	}
	return items, rows.Err()
}

func (s *FileService) ListSystem(ctx context.Context) ([]FileObject, error) {
	rows, err := s.assets.pool.Query(ctx, fileObjectSelect+` WHERE scope='system' ORDER BY created_at DESC`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	items := []FileObject{}
	for rows.Next() {
		obj, err := scanFileObject(rows)
		if err != nil {
			return nil, err
		}
		items = append(items, obj)
	}
	return items, rows.Err()
}

func (s *FileService) GetPersonal(ctx context.Context, ownerUserID, id string) (FileObject, error) {
	return scanFileObject(s.assets.pool.QueryRow(ctx, fileObjectSelect+` WHERE id=$1 AND scope='personal' AND owner_user_id=$2`, id, ownerUserID))
}

func (s *FileService) GetSystem(ctx context.Context, id string) (FileObject, error) {
	return scanFileObject(s.assets.pool.QueryRow(ctx, fileObjectSelect+` WHERE id=$1 AND scope='system'`, id))
}

// Open streams the object body for download.
func (s *FileService) Open(ctx context.Context, obj FileObject) (io.ReadCloser, int64, error) {
	return s.store.OpenObject(ctx, obj.ObjectKey)
}

func (s *FileService) DeletePersonal(ctx context.Context, ownerUserID, id string) error {
	return s.deleteFile(ctx, id, `scope='personal' AND owner_user_id=$2`, ownerUserID)
}

func (s *FileService) DeleteSystem(ctx context.Context, id string) error {
	return s.deleteFile(ctx, id, `scope='system'`, nil)
}

// deleteFile removes the metadata row and, best-effort, the object body. The
// object removal is best-effort because an orphan in the bucket is recoverable
// and must never block the delete that the user asked for.
func (s *FileService) deleteFile(ctx context.Context, id, where string, owner any) error {
	// 只传 SQL 中真实出现的占位符参数:system 范围没有 $2,多传 nil 会让
	// Postgres 报 "bind message supplies 2 parameters" 错误。
	args := []any{id}
	if owner != nil {
		args = append(args, owner)
	}
	var key string
	if err := s.assets.pool.QueryRow(ctx, `SELECT object_key FROM gateway.file_objects WHERE id=$1 AND `+where, args...).Scan(&key); err != nil {
		return mapNotFound(err)
	}
	tag, err := s.assets.pool.Exec(ctx, `DELETE FROM gateway.file_objects WHERE id=$1 AND `+where, args...)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return ErrNotFound
	}
	_ = s.store.DeleteObject(ctx, key)
	return nil
}